| from modal import App, Image, Volume |
| import modal |
| import json |
| model_name = "Learnable" |
| app = App(f"Train Model {model_name} with Text8") |
| |
| image = ( |
| Image.from_registry("nvidia/cuda:12.4.0-devel-ubuntu22.04",add_python="3.10") |
| .pip_install(["torch", "flax", "pandas", "tqdm","optax", "dataclasses", "argparse","matplotlib", |
| "scikit-learn","wandb","timm","torchvision","datasets","transformers","timm"]) |
| .run_commands("""pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html""") |
| .add_local_file("pyproject.toml", "/root/pyproject.toml") |
| .add_local_dir("src", "/root/src") |
| ) |
| |
| volume1 = Volume.from_name("weights", create_if_missing=True) |
| volume2 = Volume.from_name("datasets", create_if_missing=True) |
| volume3 = Volume.from_name("plots", create_if_missing=True) |
| volume4 = Volume.from_name("results", create_if_missing=True) |
| @app.function( |
| image=image, |
| gpu="H100:1", |
| timeout=3600 * 24, |
| volumes={"/root/weights/":volume1,"/root/datasets/":volume2,"/root/plots/":volume3,"/root/results/":volume4}, |
| ) |
| def run_command(): |
| import os |
| os.system("pip install -e .") |
| os.environ["NVIDIA_TF32_OVERRIDE"] = "0" |
| os.environ["JAX_DEFAULT_MATMUL_PRECISION"] = "highest" |
| os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" |
| os.system(""" |
| WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 python src/lgmodeling/finetune.py \ |
| --model-path /root/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \ |
| --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 8 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ |
| --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset text8 \ |
| --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\ |
| --model-save-dir /root/weights/text8/gpt2-finetune --data-path /root/datasets/text8 |
| """) |
| os.system(""" |
| WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 python src/lgmodeling/finetune.py \ |
| --model-path /root/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \ |
| --seed 20 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 8 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ |
| --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset text8 \ |
| --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\ |
| --model-save-dir /root/weights/text8/gpt2-finetune --data-path /root/datasets/text8 |
| """) |
| os.system(""" |
| WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 python src/lgmodeling/finetune.py \ |
| --model-path /root/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \ |
| --seed 40 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 8 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ |
| --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset text8 \ |
| --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\ |
| --model-save-dir /root/weights/text8/gpt2-finetune --data-path /root/datasets/text8 |
| """) |
| os.system(""" |
| WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 python src/lgmodeling/finetune.py \ |
| --model-path /root/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \ |
| --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 16 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ |
| --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset text8 \ |
| --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\ |
| --model-save-dir /root/weights/text8/gpt2-finetune --data-path /root/datasets/text8 |
| """) |
| os.system(""" |
| WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 python src/lgmodeling/finetune.py \ |
| --model-path /root/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \ |
| --seed 20 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 16 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ |
| --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset text8 \ |
| --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\ |
| --model-save-dir /root/weights/text8/gpt2-finetune --data-path /root/datasets/text8 |
| """) |
| os.system(""" |
| WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 python src/lgmodeling/finetune.py \ |
| --model-path /root/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \ |
| --seed 40 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 16 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ |
| --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset text8 \ |
| --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\ |
| --model-save-dir /root/weights/text8/gpt2-finetune --data-path /root/datasets/text8 |
| """) |
| os.system(""" |
| WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 python src/lgmodeling/finetune.py \ |
| --model-path /root/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \ |
| --seed 0 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 4 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ |
| --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset text8 \ |
| --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\ |
| --model-save-dir /root/weights/text8/gpt2-finetune --data-path /root/datasets/text8 |
| """) |
| os.system(""" |
| WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 python src/lgmodeling/finetune.py \ |
| --model-path /root/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \ |
| --seed 20 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 4 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ |
| --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset text8 \ |
| --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\ |
| --model-save-dir /root/weights/text8/gpt2-finetune --data-path /root/datasets/text8 |
| """) |
| os.system(""" |
| WANDB_MODE=online CUDA_VISIBLE_DEVICES=0 python src/lgmodeling/finetune.py \ |
| --model-path /root/weights/text8/lr0.00025-learnable-step60000-warm0-size24-layer12-embd512-heads8-shared1-routed0-topk0/best_60000 \ |
| --seed 40 --tgt_len 512 --mem_len 512 --eval_tgt_len 128 --n_head 4 --lmc-layer-indices 0 1 2 3 4 5 6 7 8 9 10 11\ |
| --learning-rate 0.00025 --batch-size 24 --max_step 60000 --warmup_step 0 --dataset text8 \ |
| --wandb-project LMC-Attention --wandb-group "GPT2-Text8-FFN" --wandb-entity "vinh-bui0512-hcmut"\ |
| --model-save-dir /root/weights/text8/gpt2-finetune --data-path /root/datasets/text8 |
| """) |
|
|
| if __name__ == "__main__": |
| with app.run(): |
| run_command.remote() |
