diff --git a/README.md b/README.md index 5f7c625..ae506a1 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,7 @@ Use this repository to learn how to build distributed training jobs on Modal. - [**`nanoGPT/`**](/nanoGPT/) training Karpathy's nanoGPT reproduction of OpenAI's GPT-2. - [**`resnet50/`**](/resnet50/) training a ResNet50 model on the ImageNet dataset. - [**`starcoder/`**](/starcoder) accelerated finetuning of Llama-2-7B on Rust and Go code, supporting either `torchrun` or `accelerate`. +- [**`ms-swift/`**](/ms-swift/) LoRA SFT of GLM-4.7 with ms-swift on the same LongMIT-128K dataset used in the Megatron example. # Documentation diff --git a/ms-swift/README.md b/ms-swift/README.md new file mode 100644 index 0000000..8a08c35 --- /dev/null +++ b/ms-swift/README.md @@ -0,0 +1,73 @@ +# ms-swift: SFT GLM-4.7 (same LongMIT dataset as Megatron example) + +This example shows supervised fine-tuning (SFT) of `zai-org/GLM-4.7` with +[ms-swift](https://github.com/modelscope/ms-swift) on Modal. + +It intentionally reuses the same dataset source and preprocessing strategy as +the Megatron example: + +- Dataset source: `donmaclean/LongMIT-128K` +- Prompt construction: passages + question -> answer +- Token filter: `<= 131072` +- Output path/schema: `/data/longmit-128k/training.jsonl` with `input`/`output` + +## Prerequisites + +- Modal account with H100 access +- Hugging Face token with access to GLM-4.7 + +Create the required secret: + +```bash +modal secret create huggingface-secret HF_TOKEN= +``` + +Optional (only if you want W&B logging): pass `WANDB_API_KEY` as an env var +when launching training, for example: + +```bash +WANDB_API_KEY= modal run --detach ms-swift/modal_train.py::train_sft +``` + +## Quick Start + +### 1) Download model weights + +```bash +modal run ms-swift/modal_train.py::download_model +``` + +### 2) Prepare dataset (same as Megatron example) + +```bash +modal run ms-swift/modal_train.py::prep_dataset +``` + +### 3) Launch multi-node SFT with ms-swift + +```bash +modal run --detach ms-swift/modal_train.py::train_sft +``` + +## What `train_sft` uses by default + +- 4 nodes x 8 H100 GPUs +- `swift sft` + LoRA +- DeepSpeed ZeRO-3 +- BF16 +- `max_length=16384` with `truncation_strategy=left` +- same preprocessed LongMIT JSONL from `/data/longmit-128k/training.jsonl` + +## Customize + +You can override key args directly from `modal run`, for example: + +```bash +modal run --detach ms-swift/modal_train.py::train_sft \ + --run-name glm47-ms-swift-longrun \ + --num-train-epochs 2 \ + --per-device-train-batch-size 1 \ + --gradient-accumulation-steps 8 \ + --max-length 32768 +``` + diff --git a/ms-swift/modal_train.py b/ms-swift/modal_train.py new file mode 100644 index 0000000..7372b68 --- /dev/null +++ b/ms-swift/modal_train.py @@ -0,0 +1,289 @@ +""" +GLM-4.7 SFT with ms-swift on Modal. + +Usage: + modal run ms-swift/modal_train.py::download_model + modal run ms-swift/modal_train.py::prep_dataset + modal run --detach ms-swift/modal_train.py::train_sft +""" + +import json +import os +from typing import Any, Dict, Tuple + +import modal +import modal.experimental + +app = modal.App("glm47-ms-swift-sft") + +# Volumes +models_volume = modal.Volume.from_name("big-model-hfcache", create_if_missing=True) +data_volume = modal.Volume.from_name("glm47-training-data", create_if_missing=True) +checkpoints_volume = modal.Volume.from_name("glm47-ms-swift-checkpoints", create_if_missing=True) + +# Paths +HF_CACHE = "/root/.cache/huggingface" +DATA_DIR = "/data" +CHECKPOINTS_DIR = "/checkpoints" +PREPROCESSED_DIR = f"{DATA_DIR}/longmit-128k" +TRAIN_JSONL = f"{PREPROCESSED_DIR}/training.jsonl" + +# Models/Datasets +HF_MODEL = "zai-org/GLM-4.7" +HF_DATASET = "donmaclean/LongMIT-128K" + +# Training/Dataset defaults +MAX_SFT_TOKENS = 131_072 +DEFAULT_MAX_LENGTH = 16_384 +PREP_CPU = 32 + +# Distributed training defaults +N_NODES = 4 +GPUS_PER_NODE = 8 +MASTER_PORT = 29500 + +download_image = ( + modal.Image.debian_slim(python_version="3.11") + .uv_pip_install( + "huggingface_hub==0.36.0", + "transformers==4.57.4", + "torch==2.9.1", + "safetensors==0.7.0", + "sentencepiece==0.2.1", + ) + .env({"HF_XET_HIGH_PERFORMANCE": "1"}) +) + +prep_image = download_image.uv_pip_install("datasets==3.1.0") + +swift_image = ( + modal.Image.from_registry("nvcr.io/nvidia/nemo:25.11") + .entrypoint([]) + .uv_pip_install( + "ms-swift==3.12.4", + "deepspeed==0.17.6", + ) + .run_commands(f"rm -Rf {HF_CACHE}") +) + + +def _build_longmit_prompt(example: Dict[str, Any]) -> Tuple[str, str]: + passages = "\n".join( + [f"Passage {i + 1}:\n{doc['content']}" for i, doc in enumerate(example["all_docs"])] + ) + prompt = ( + "Answer the question based on the given passages.\n\n" + f"{passages}\n\n" + f"Question: {example['question']}\nAnswer:" + ) + return prompt, example["answer"] + + +@app.function( + image=download_image, + volumes={HF_CACHE: models_volume}, + secrets=[modal.Secret.from_name("huggingface-secret")], + timeout=14400, +) +def download_model(): + """Download GLM-4.7 weights to the HuggingFace cache volume.""" + from huggingface_hub import snapshot_download + + models_volume.reload() + + print(f"Downloading model: {HF_MODEL}") + path = snapshot_download(HF_MODEL, token=os.environ.get("HF_TOKEN")) + print(f"Model downloaded to: {path}") + + models_volume.commit() + return {"model_path": path} + + +@app.function( + image=prep_image, + volumes={DATA_DIR: data_volume, HF_CACHE: models_volume}, + secrets=[modal.Secret.from_name("huggingface-secret")], + timeout=7200, + cpu=PREP_CPU, +) +def prep_dataset(): + """ + Prepare LongMIT-128K exactly like the Megatron example: + - source dataset: donmaclean/LongMIT-128K + - same prompt format + - same <=131072 token filtering + - same output path: /data/longmit-128k/training.jsonl + """ + from datasets import load_dataset + from transformers import AutoTokenizer + + data_volume.reload() + models_volume.reload() + os.makedirs(PREPROCESSED_DIR, exist_ok=True) + + print(f"Loading dataset from Hugging Face: {HF_DATASET}") + dataset = load_dataset(HF_DATASET, split="train", trust_remote_code=True) + print(f"Loaded {len(dataset)} examples") + + tokenizer = AutoTokenizer.from_pretrained( + HF_MODEL, + use_fast=True, + trust_remote_code=True, + ) + + def format_longmit(example: Dict[str, Any]) -> Dict[str, Any]: + prompt, answer = _build_longmit_prompt(example) + token_count = len(tokenizer(prompt + answer).input_ids) + return {"input": prompt, "output": answer, "n_tokens": token_count} + + print(f"Formatting and filtering to <= {MAX_SFT_TOKENS} tokens") + original_len = len(dataset) + dataset = dataset.map( + format_longmit, + remove_columns=dataset.column_names, + num_proc=PREP_CPU, + ) + dataset = dataset.filter( + lambda ex: ex["n_tokens"] <= MAX_SFT_TOKENS, + num_proc=PREP_CPU, + ) + filtered_len = len(dataset) + print(f"Filtered dataset size: {filtered_len} (from {original_len})") + + print(f"Writing JSONL: {TRAIN_JSONL}") + with open(TRAIN_JSONL, "w", encoding="utf-8") as f: + for example in dataset: + # Keep input/output schema aligned with the Megatron example. + json.dump({"input": example["input"], "output": example["output"]}, f) + f.write("\n") + + fsize = os.path.getsize(TRAIN_JSONL) + print(f"Created {TRAIN_JSONL} ({fsize:,} bytes)") + + data_volume.commit() + return {"dataset_path": TRAIN_JSONL, "examples": filtered_len} + + +@app.function( + image=swift_image, + gpu="H100:8", + volumes={ + HF_CACHE: models_volume, + DATA_DIR: data_volume, + CHECKPOINTS_DIR: checkpoints_volume, + }, + secrets=[modal.Secret.from_name("huggingface-secret")], + timeout=86400, + experimental_options={"efa_enabled": True}, +) +@modal.experimental.clustered(size=N_NODES, rdma=True) +def train_sft( + run_name: str = "glm47-ms-swift-lora", + num_train_epochs: int = 1, + per_device_train_batch_size: int = 1, + gradient_accumulation_steps: int = 4, + max_length: int = DEFAULT_MAX_LENGTH, +): + """ + Multi-node SFT for GLM-4.7 with ms-swift. + + This uses the same preprocessed dataset path as the Megatron example: + /data/longmit-128k/training.jsonl + """ + import subprocess + + data_volume.reload() + models_volume.reload() + + if not os.path.exists(TRAIN_JSONL): + raise RuntimeError( + f"Dataset missing at {TRAIN_JSONL}. " + "Run `modal run ms-swift/modal_train.py::prep_dataset` first." + ) + + output_dir = f"{CHECKPOINTS_DIR}/{run_name}" + + cluster_info = modal.experimental.get_cluster_info() + node_rank = cluster_info.rank + master_addr = cluster_info.container_ips[0] if cluster_info.container_ips else "localhost" + + # `swift` auto-switches to torchrun when these env vars are present. + os.environ["NNODES"] = str(N_NODES) + os.environ["NPROC_PER_NODE"] = str(GPUS_PER_NODE) + os.environ["NODE_RANK"] = str(node_rank) + os.environ["MASTER_ADDR"] = master_addr + os.environ["MASTER_PORT"] = str(MASTER_PORT) + + cmd = [ + "swift", + "sft", + "--model", + HF_MODEL, + "--dataset", + TRAIN_JSONL, + "--use_hf", + "true", + "--train_type", + "lora", + "--deepspeed", + "zero3", + "--torch_dtype", + "bfloat16", + "--target_modules", + "all-linear", + "--lora_rank", + "128", + "--lora_alpha", + "32", + "--lora_dropout", + "0.05", + "--num_train_epochs", + str(num_train_epochs), + "--per_device_train_batch_size", + str(per_device_train_batch_size), + "--gradient_accumulation_steps", + str(gradient_accumulation_steps), + "--learning_rate", + "1e-4", + "--save_strategy", + "steps", + "--save_steps", + "50", + "--save_total_limit", + "2", + "--logging_steps", + "5", + "--max_length", + str(max_length), + "--truncation_strategy", + "left", + "--split_dataset_ratio", + "0", + "--dataloader_num_workers", + "4", + "--gradient_checkpointing", + "true", + "--output_dir", + output_dir, + ] + + hf_token = os.environ.get("HF_TOKEN") + if hf_token: + cmd.extend(["--hub_token", hf_token]) + + wandb_key = os.environ.get("WANDB_API_KEY") + if wandb_key: + cmd.extend(["--report_to", "wandb", "--run_name", run_name]) + else: + cmd.extend(["--report_to", "tensorboard"]) + + print(f"Starting rank {node_rank}/{N_NODES - 1} with master {master_addr}:{MASTER_PORT}") + print("Running command:") + print(" ".join(cmd)) + + result = subprocess.run(cmd, check=False) + if result.returncode != 0: + raise RuntimeError(f"swift sft exited with code {result.returncode}") + + checkpoints_volume.commit() + return {"status": "complete", "output_dir": output_dir}