diff --git a/README.md b/README.md index 2f32fc9..1993482 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,7 @@ Use this repository to learn how to build distributed training jobs on Modal. # Examples - [**`benchmark/`**](/benchmark/) contains performance and reliability testing, using AWS EFA by default. +- [**`deepseek-v4-flash-sft/`**](/deepseek-v4-flash-sft/) 60k-context LoRA SFT and serving for DeepSeek-V4-Flash using NeMo AutoModel and vLLM. - [**`lightning/`**](/lightning/) a simple lightning.ai Fabric example. - [**`nanoGPT/`**](/nanoGPT/) training Karpathy's nanoGPT reproduction of OpenAI's GPT-2. - [**`resnet50/`**](/resnet50/) training a ResNet50 model on the ImageNet dataset. diff --git a/deepseek-v4-flash-sft/README.md b/deepseek-v4-flash-sft/README.md new file mode 100644 index 0000000..1d706d2 --- /dev/null +++ b/deepseek-v4-flash-sft/README.md @@ -0,0 +1,130 @@ +# DeepSeek-V4-Flash 60k LoRA SFT and serving + +Train a rank-64 attention LoRA for +[`deepseek-ai/DeepSeek-V4-Flash`](https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash) +at a 60,000-token sequence length with NeMo AutoModel, finalize the +pipeline-parallel checkpoint as a PEFT adapter, and serve it with vLLM. + +## Requirements + +1. Modal multi-node cluster access and quota for 8 nodes of 8 H200 GPUs. +2. A Modal secret named `huggingface-secret` containing `HF_TOKEN`. +3. The repository's locked `uv` environment. + +Create the Hugging Face secret once: + +```bash +uv run --frozen modal secret create huggingface-secret HF_TOKEN=hf_xxxxx +``` + +## Topology + +| Dimension | 8-node default | 16-node option | +| --- | ---: | ---: | +| Nodes / GPUs | 8 / 64 H200 | 16 / 128 H200 | +| TP / DP / PP / CP / EP | 1 / 1 / 4 / 16 / 16 | 1 / 2 / 4 / 16 / 32 | +| Sequence length | 60,000 | 60,000 | +| Global / local batch | 8 / 4 | 8 / 4 | +| Gradient accumulation | 2 | 1 | +| Attention / MoE dispatch | TileLang / UCCL-EP | TileLang / UCCL-EP | +| LoRA | rank 64, alpha 64 | rank 64, alpha 64 | +| AutoModel targets | `wq_a`, `wq_b`, `wkv` | `wq_a`, `wq_b`, `wkv` | +| Exported PEFT targets | `q_a_proj`, `q_b_proj`, `kv_proj` | `q_a_proj`, `q_b_proj`, `kv_proj` | +| Optimizer steps | 5 | 5 | +| Checkpoint step | 4 (zero-based) | 4 (zero-based) | +| Scheduler / cgroup memory | 128 MiB / 256 GiB per node | 128 MiB / 256 GiB per node | + +The request enables EFA so Modal can schedule from the larger RDMA-capable +pool. The image contains UCCL extensions for both EFA and Mellanox/RoCE and +activates the extension that matches the scheduled nodes. + +## Train + +Run from the repository root with a unique run ID: + +```bash +uv run --frozen modal run --detach \ + deepseek-v4-flash-sft/modal_train.py::train \ + --run-id dsv4-flash-60k-lora-$(date -u +%Y%m%d-%H%M%S) +``` + +Eight nodes is the default. Set `N_NODES=16` before the command to select the +DP=2, EP=32 layout. The launcher derives the topology from `N_NODES` and +rejects incompatible layouts before scheduling. + +The training function writes a fixed-length synthetic chat dataset on each +node and renders [`train_recipe.yaml`](train_recipe.yaml) with the selected EP +size and checkpoint path. Replace `_write_synthetic_dataset` and the recipe's +`dataset` section to use a production dataset. + +Each pipeline stage writes its adapter shard to the +`example-deepseek-v4-flash-sft-checkpoints` volume. The run ID is immutable: +the trainer refuses to overwrite an existing checkpoint directory. + +## Finalize + +After the detached training app finishes, merge the four pipeline-stage shards: + +```bash +uv run --frozen modal run \ + deepseek-v4-flash-sft/modal_train.py::finalize \ + --run-id +``` + +Finalization requires all 258 tensors across all 43 layers, verifies their +shapes and values, rewrites AutoModel projection names to the Transformers +names used by PEFT and vLLM, and loads the complete adapter through PEFT. + +The finalized files are: + +```text +/checkpoints//epoch_0_step_4/ +|-- checkpoint_manifest.json +`-- model/ + |-- adapter_config.json + `-- adapter_model.safetensors +``` + +The manifest records the adapter SHA-256 digest. Optimizer state is not saved, +so this output is a serving adapter rather than a training-resume checkpoint. + +## Serve + +Deploy the finalized run: + +```bash +SERVE_RUN_ID= \ + uv run --frozen modal deploy \ + deepseek-v4-flash-sft/modal_train.py +``` + +The `serve` endpoint is an OpenAI-compatible vLLM server protected by Modal +proxy authentication. It exposes the base model as +`deepseek-ai/DeepSeek-V4-Flash` and the adapter as +`deepseek-v4-flash-60k-lora`, with a maximum model length of 65,536 tokens. + +Serving uses pipeline parallelism across four H200 GPUs. The launcher verifies +the finalized manifest and adapter digest before starting vLLM. + +## Integration patches + +The pinned AutoModel, UCCL, and vLLM images use six integration patches: + +| Patch | Scope | Purpose | +| --- | --- | --- | +| [`patches/automodel_checkpoint_dequant.patch`](patches/automodel_checkpoint_dequant.patch) | AutoModel loading | Pass FP4 checkpoint dequantization through the recipe API. | +| [`patches/automodel_pp_peft_checkpoint.patch`](patches/automodel_pp_peft_checkpoint.patch) | AutoModel checkpointing | Save one PEFT shard per pipeline stage without optimizer state. | +| [`patches/automodel_composite_backend.patch`](patches/automodel_composite_backend.patch) | AutoModel control plane | Handle the `cpu:gloo,cuda:nccl` backend in signal handling. | +| [`patches/automodel_uccl_teardown.patch`](patches/automodel_uccl_teardown.patch) | AutoModel shutdown | Release UCCL before distributed and CUDA teardown. | +| [`patches/uccl_ipv6_oob.patch`](patches/uccl_ipv6_oob.patch) | UCCL transport | Use Modal's IPv6 inter-node interface for OOB coordination. | +| [`patches/vllm_deepseek_v4_lora.patch`](patches/vllm_deepseek_v4_lora.patch) | vLLM adapter loading | Register DeepSeek-V4 LoRA and map PEFT attention names to fused modules. | + +These patches cover model loading, checkpoint export, transport, shutdown, and +adapter registration. They do not modify the model forward pass, attention +math, or autograd graph. + +vLLM 0.25.1 provides DeepSeek-V4 serving and pipeline parallelism, while the +patch supplies the LoRA registration and packed Q/KV mapping for +`DeepseekV4ForCausalLM`. The serving image downloads the official model config +to an isolated directory so the FP8 quantization metadata comes directly from +the model repository. diff --git a/deepseek-v4-flash-sft/dsv4_checkpoint.py b/deepseek-v4-flash-sft/dsv4_checkpoint.py new file mode 100644 index 0000000..fa18879 --- /dev/null +++ b/deepseek-v4-flash-sft/dsv4_checkpoint.py @@ -0,0 +1,316 @@ +# pyright: reportMissingImports=false +"""Finalize and verify DeepSeek-V4-Flash pipeline-parallel LoRA checkpoints.""" + +from __future__ import annotations + +import hashlib +import json +import re +from pathlib import Path + +HF_MODEL = "deepseek-ai/DeepSeek-V4-Flash" +MODEL_LAYERS = 43 +LORA_RANK = 64 +LORA_ALPHA = 64 +PIPELINE_STAGES = 4 +CHECKPOINT_STEP = 4 + +TARGET_MAP = { + "wq_a": "q_a_proj", + "wq_b": "q_b_proj", + "wkv": "kv_proj", +} +TARGET_SHAPES = { + "wq_a": { + "lora_A": (LORA_RANK, 4096), + "lora_B": (1024, LORA_RANK), + }, + "wq_b": { + "lora_A": (LORA_RANK, 1024), + "lora_B": (32768, LORA_RANK), + }, + "wkv": { + "lora_A": (LORA_RANK, 4096), + "lora_B": (512, LORA_RANK), + }, +} +PEFT_TARGET_PATTERN = r"model\.layers\.\d+\.self_attn\.(q_a_proj|q_b_proj|kv_proj)" +EXPECTED_TENSOR_COUNT = MODEL_LAYERS * len(TARGET_MAP) * 2 + + +def _validate_run_id(run_id: str) -> None: + if not run_id or Path(run_id).name != run_id: + raise ValueError("run_id must be a non-empty path component") + + +def finalized_adapter( + checkpoints_dir: str | Path, + run_id: str, +) -> tuple[Path, dict[str, object]]: + """Return a finalized adapter after verifying its manifest and digest.""" + _validate_run_id(run_id) + checkpoint_dir = Path(checkpoints_dir) / run_id / f"epoch_0_step_{CHECKPOINT_STEP}" + model_dir = checkpoint_dir / "model" + adapter_path = model_dir / "adapter_model.safetensors" + config_path = model_dir / "adapter_config.json" + manifest_path = checkpoint_dir / "checkpoint_manifest.json" + for path in (adapter_path, config_path, manifest_path): + if not path.is_file(): + raise FileNotFoundError(f"Missing finalized adapter artifact: {path}") + + manifest = json.loads(manifest_path.read_text()) + if ( + manifest.get("run_id") != run_id + or manifest.get("step") != CHECKPOINT_STEP + or manifest.get("adapter_tensor_count") != EXPECTED_TENSOR_COUNT + ): + raise RuntimeError( + f"Checkpoint manifest does not match run {run_id!r}: {manifest}" + ) + + digest = hashlib.sha256(adapter_path.read_bytes()).hexdigest() + if digest != manifest.get("adapter_sha256"): + raise RuntimeError( + f"Adapter SHA-256 {digest} does not match {manifest.get('adapter_sha256')}" + ) + + adapter_config = json.loads(config_path.read_text()) + if ( + adapter_config.get("r") != LORA_RANK + or adapter_config.get("lora_alpha") != LORA_ALPHA + ): + raise RuntimeError( + f"Expected rank-{LORA_RANK}/alpha-{LORA_ALPHA}, found " + f"r={adapter_config.get('r')}, " + f"alpha={adapter_config.get('lora_alpha')}" + ) + return model_dir, manifest + + +def _validate_with_peft(model_dir: Path) -> dict[str, object]: + """Load every adapter tensor through PEFT without materializing the base.""" + import torch + from accelerate import init_empty_weights + from peft import PeftConfig, get_peft_model + from peft.utils.save_and_load import ( + load_peft_weights, + set_peft_model_state_dict, + ) + from transformers import AutoConfig, AutoModelForCausalLM + + base_config = AutoConfig.from_pretrained(HF_MODEL, trust_remote_code=False) + adapter_config = PeftConfig.from_pretrained(str(model_dir)) + with init_empty_weights(include_buffers=True): + base_model = AutoModelForCausalLM.from_config( + base_config, + trust_remote_code=False, + dtype=torch.bfloat16, + ) + model = get_peft_model(base_model, adapter_config, low_cpu_mem_usage=True) + adapter_params = [ + (name, param) for name, param in model.named_parameters() if "lora_" in name + ] + if len(adapter_params) != EXPECTED_TENSOR_COUNT: + raise RuntimeError( + f"PEFT injected {len(adapter_params)} adapter parameters, expected " + f"{EXPECTED_TENSOR_COUNT}" + ) + + for module in model.modules(): + if hasattr(module, "lora_A") and "default" in module.lora_A: + module.lora_A["default"].to_empty(device="cpu") + module.lora_B["default"].to_empty(device="cpu") + + state = load_peft_weights(str(model_dir), device="cpu") + load_result = set_peft_model_state_dict(model, state, adapter_name="default") + missing = [key for key in load_result.missing_keys if "lora_" in key] + if missing: + raise RuntimeError(f"PEFT load is missing adapter keys: {missing[:10]}") + if load_result.unexpected_keys: + raise RuntimeError( + f"PEFT load has unexpected keys: {load_result.unexpected_keys[:10]}" + ) + + adapter_params = [ + (name, param) for name, param in model.named_parameters() if "lora_" in name + ] + invalid = [ + name + for name, param in adapter_params + if param.is_meta + or not torch.isfinite(param).all() + or (".lora_B." in name and torch.count_nonzero(param).item() == 0) + ] + if invalid: + raise RuntimeError(f"PEFT loaded invalid adapter parameters: {invalid[:10]}") + + layers = sorted( + { + int(match.group(1)) + for name, _ in adapter_params + if (match := re.search(r"\.layers\.(\d+)\.", name)) + } + ) + targets = sorted( + { + match.group(1) + for name, _ in adapter_params + if ( + match := re.search( + r"\.self_attn\.(q_a_proj|q_b_proj|kv_proj)\.", + name, + ) + ) + } + ) + if layers != list(range(MODEL_LAYERS)): + raise RuntimeError(f"PEFT load covered unexpected layers: {layers}") + if targets != ["kv_proj", "q_a_proj", "q_b_proj"]: + raise RuntimeError(f"PEFT load covered unexpected targets: {targets}") + + return { + "adapter_parameters": len(adapter_params), + "layers": len(layers), + "targets": targets, + } + + +def finalize_adapter( + checkpoints_dir: str | Path, + run_id: str, +) -> dict[str, object]: + """Merge PP-local shards into one PEFT adapter and verify it.""" + import torch + from safetensors.torch import load_file, save_file + + _validate_run_id(run_id) + run_dir = Path(checkpoints_dir) / run_id + candidates: list[tuple[tuple[int, int], Path]] = [] + checkpoint_pattern = re.compile(r"epoch_(\d+)_step_(\d+)") + if run_dir.exists(): + for path in run_dir.iterdir(): + match = checkpoint_pattern.fullmatch(path.name) + if path.is_dir() and match: + candidates.append(((int(match.group(1)), int(match.group(2))), path)) + if not candidates: + raise FileNotFoundError(f"No checkpoint steps found under {run_dir}") + + (epoch, step), checkpoint_dir = max(candidates) + if step != CHECKPOINT_STEP: + raise RuntimeError( + f"Latest checkpoint is step {step}, expected {CHECKPOINT_STEP}" + ) + + model_dir = checkpoint_dir / "model" + stage_shards = sorted(model_dir.glob("adapter_model.pp-*-of-*.safetensors")) + if len(stage_shards) != PIPELINE_STAGES: + raise RuntimeError( + f"Found {len(stage_shards)} PP shards, expected {PIPELINE_STAGES}" + ) + + merged: dict[str, torch.Tensor] = {} + for shard_path in stage_shards: + shard = load_file(str(shard_path), device="cpu") + duplicate_keys = sorted(set(merged).intersection(shard)) + if duplicate_keys: + raise RuntimeError( + f"Duplicate keys in {shard_path.name}: {duplicate_keys[:10]}" + ) + nonfinite = [ + name for name, tensor in shard.items() if not torch.isfinite(tensor).all() + ] + if nonfinite: + raise RuntimeError( + f"Non-finite tensors in {shard_path.name}: {nonfinite[:10]}" + ) + merged.update(shard) + + if len(merged) != EXPECTED_TENSOR_COUNT: + raise RuntimeError( + f"Merged adapter has {len(merged)} tensors, " + f"expected {EXPECTED_TENSOR_COUNT}" + ) + + expected_adapters = ("lora_A", "lora_B") + matched_names: set[str] = set() + peft_state: dict[str, torch.Tensor] = {} + for layer in range(MODEL_LAYERS): + for source_target, peft_target in TARGET_MAP.items(): + for adapter in expected_adapters: + needle = f".layers.{layer}.self_attn.{source_target}.{adapter}.weight" + matches = [key for key in merged if needle in key] + if len(matches) != 1: + raise RuntimeError( + f"Expected one tensor matching {needle}, found {len(matches)}" + ) + + name = matches[0] + tensor = merged[name] + expected_shape = TARGET_SHAPES[source_target][adapter] + if tuple(tensor.shape) != expected_shape: + raise RuntimeError( + f"{name} has shape {tuple(tensor.shape)}, " + f"expected {expected_shape}" + ) + if adapter == "lora_B" and torch.count_nonzero(tensor).item() == 0: + raise RuntimeError(f"Untrained LoRA-B tensor: {name}") + + matched_names.add(name) + peft_name = name.replace( + f".self_attn.{source_target}.", + f".self_attn.{peft_target}.", + 1, + ) + if peft_name in peft_state: + raise RuntimeError(f"Duplicate PEFT key: {peft_name}") + peft_state[peft_name] = tensor + + unexpected = sorted(set(merged).difference(matched_names)) + if unexpected: + raise RuntimeError(f"Unexpected adapter tensors: {unexpected[:10]}") + + adapter_config_path = model_dir / "adapter_config.json" + if not adapter_config_path.is_file(): + raise FileNotFoundError(f"Missing PEFT config: {adapter_config_path}") + adapter_config = json.loads(adapter_config_path.read_text()) + if ( + adapter_config.get("r") != LORA_RANK + or adapter_config.get("lora_alpha") != LORA_ALPHA + ): + raise RuntimeError( + f"Expected rank-{LORA_RANK}/alpha-{LORA_ALPHA}, found " + f"r={adapter_config.get('r')}, " + f"alpha={adapter_config.get('lora_alpha')}" + ) + adapter_config["target_modules"] = PEFT_TARGET_PATTERN + + temporary_config = adapter_config_path.with_suffix(".json.tmp") + temporary_config.write_text( + json.dumps(adapter_config, indent=2, sort_keys=True) + "\n" + ) + temporary_config.replace(adapter_config_path) + + adapter_path = model_dir / "adapter_model.safetensors" + temporary_adapter = adapter_path.with_suffix(".safetensors.tmp") + save_file(peft_state, str(temporary_adapter)) + temporary_adapter.replace(adapter_path) + + peft_validation = _validate_with_peft(model_dir) + manifest = { + "run_id": run_id, + "epoch": epoch, + "step": step, + "pipeline_stages": PIPELINE_STAGES, + "adapter_tensor_count": len(peft_state), + "adapter_bytes": adapter_path.stat().st_size, + "adapter_sha256": hashlib.sha256(adapter_path.read_bytes()).hexdigest(), + "layers": MODEL_LAYERS, + "lora_rank": LORA_RANK, + "lora_alpha": LORA_ALPHA, + "targets": list(TARGET_MAP.values()), + "peft_validation": peft_validation, + } + (checkpoint_dir / "checkpoint_manifest.json").write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n" + ) + return manifest diff --git a/deepseek-v4-flash-sft/dsv4_images.py b/deepseek-v4-flash-sft/dsv4_images.py new file mode 100644 index 0000000..0cd9913 --- /dev/null +++ b/deepseek-v4-flash-sft/dsv4_images.py @@ -0,0 +1,174 @@ +"""Container images for DeepSeek-V4-Flash training and serving.""" + +from pathlib import Path + +import modal + +AUTOMODEL_IMAGE = "nvcr.io/nvidia/nemo-automodel:26.06.00" +AUTOMODEL_COMMIT = "1197b6281255957cc2c58e79d50796c1b256c57c" +UCCL_EP_COMMIT = "66170bc299205228f0170bc1638594a39af9ffd5" +UCCL_EFA_INSTALLER_VERSION = "1.42.0" +VLLM_IMAGE = "vllm/vllm-openai:v0.25.1" + +REMOTE_TRAIN_RECIPE = "/root/train_recipe.yaml" +REMOTE_CHECKPOINT_HELPER = "/root/dsv4_checkpoint.py" + + +def build_automodel_image(source_dir: Path) -> modal.Image: + patches_dir = source_dir / "patches" + uccl_ipv6_patch = patches_dir / "uccl_ipv6_oob.patch" + automodel_patches = tuple( + patches_dir / name + for name in ( + "automodel_uccl_teardown.patch", + "automodel_checkpoint_dequant.patch", + "automodel_composite_backend.patch", + "automodel_pp_peft_checkpoint.patch", + ) + ) + + # This commit provides the DeepSeek-V4 context-parallel implementation. + image = ( + modal.Image.from_registry(AUTOMODEL_IMAGE) + .entrypoint([]) + .run_commands( + "cd / && rm -rf /opt/Automodel && " + "git clone https://github.com/NVIDIA-NeMo/Automodel.git " + "/opt/Automodel && " + f"cd /opt/Automodel && git checkout {AUTOMODEL_COMMIT} && " + "pip install --no-deps --force-reinstall -e /opt/Automodel" + ) + .run_commands( + "pip install --no-deps --force-reinstall " + "'tilelang>=0.1.11' 'tile-kernels==1.0.0' " + "'apache-tvm-ffi<=0.1.11'" + ) + .env( + { + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "HF_XET_HIGH_PERFORMANCE": "1", + "NEMO_AUTOMODEL_DSV4_EXPERT_LAYOUT": "fp4", + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True", + "TORCH_NCCL_ASYNC_ERROR_HANDLING": "1", + } + ) + ) + + # EFA-enabled scheduling can place a job on either EFA or Mellanox/RoCE. + image = ( + image.apt_install( + "build-essential", + "curl", + "libibverbs-dev", + "libnl-3-dev", + "libnl-route-3-dev", + "libnuma-dev", + "ninja-build", + "patch", + "rdma-core", + ) + .uv_pip_install( + "intervaltree==3.1.0", + "nanobind==2.13.0", + "pybind11==3.0.1", + ) + .run_commands( + "cd /tmp && " + f"curl -fsSLO https://efa-installer.amazonaws.com/aws-efa-installer-{UCCL_EFA_INSTALLER_VERSION}.tar.gz && " + f"tar -xzf aws-efa-installer-{UCCL_EFA_INSTALLER_VERSION}.tar.gz && " + "cd aws-efa-installer && " + "./efa_installer.sh -y --skip-kmod -g --no-verify && " + "rm -rf /tmp/aws-efa-installer*" + ) + .add_local_file( + str(uccl_ipv6_patch), + "/tmp/uccl_ipv6_oob.patch", + copy=True, + ) + .run_commands( + "rm -rf /opt/uccl && mkdir -p /opt/uccl && " + f"curl -fsSL https://github.com/uccl-project/uccl/archive/{UCCL_EP_COMMIT}.tar.gz " + "| tar -xz --strip-components=1 -C /opt/uccl && " + "cd /opt/uccl && patch -p1 < /tmp/uccl_ipv6_oob.patch", + "cd /opt/uccl/ep && rm -rf build ep*.so && " + "EFA_HOME=/opt/amazon/efa USE_DMABUF=1 " + "TORCH_CUDA_ARCH_LIST=9.0 MAX_JOBS=8 " + "/opt/venv/bin/python setup.py build_ext --inplace && " + "cp ep*.so /opt/uccl/uccl/ep.efa.abi3.so", + "cd /opt/uccl/ep && rm -rf build ep*.so && " + "EFA_HOME=/opt/uccl/no-efa USE_DMABUF=1 " + "TORCH_CUDA_ARCH_LIST=9.0 MAX_JOBS=8 " + "/opt/venv/bin/python setup.py build_ext --inplace && " + "cp ep*.so /opt/uccl/uccl/ep.mellanox.abi3.so && " + "cp ep*.so /opt/uccl/uccl/ep.abi3.so", + "cd /opt/uccl && " + "uv pip install --python /opt/venv/bin/python --no-build-isolation " + "--no-deps --force-reinstall .", + ) + ) + + for patch_path in automodel_patches: + image = image.add_local_file( + str(patch_path), + f"/tmp/{patch_path.name}", + copy=True, + ) + return ( + image.run_commands( + "cd /opt/Automodel && " + + " && ".join( + f"git apply /tmp/{patch_path.name}" for patch_path in automodel_patches + ) + ) + .add_local_file( + str(source_dir / "train_recipe.yaml"), + REMOTE_TRAIN_RECIPE, + copy=True, + ) + .add_local_file( + str(source_dir / "dsv4_checkpoint.py"), + REMOTE_CHECKPOINT_HELPER, + copy=True, + ) + .add_local_file( + str(source_dir / "dsv4_images.py"), + "/root/dsv4_images.py", + copy=True, + ) + ) + + +def build_vllm_image(source_dir: Path) -> modal.Image: + lora_patch = source_dir / "patches" / "vllm_deepseek_v4_lora.patch" + return ( + modal.Image.from_registry(VLLM_IMAGE) + .entrypoint([]) + .run_commands("ln -sf $(which python3) /usr/local/bin/python") + .apt_install("patch") + .add_local_file( + str(lora_patch), + f"/tmp/{lora_patch.name}", + copy=True, + ) + .run_commands( + "cd $(python -c 'import pathlib, vllm; " + "print(pathlib.Path(vllm.__file__).resolve().parent.parent)') && " + f"patch --batch --forward -p1 < /tmp/{lora_patch.name}" + ) + .add_local_file( + str(source_dir / "dsv4_checkpoint.py"), + REMOTE_CHECKPOINT_HELPER, + copy=True, + ) + .add_local_file( + str(source_dir / "dsv4_images.py"), + "/root/dsv4_images.py", + copy=True, + ) + .env( + { + "HF_XET_HIGH_PERFORMANCE": "1", + "VLLM_ENGINE_READY_TIMEOUT_S": "3600", + } + ) + ) diff --git a/deepseek-v4-flash-sft/modal_train.py b/deepseek-v4-flash-sft/modal_train.py new file mode 100644 index 0000000..08e1c42 --- /dev/null +++ b/deepseek-v4-flash-sft/modal_train.py @@ -0,0 +1,410 @@ +# pyright: reportMissingImports=false, reportCallIssue=false, reportOptionalCall=false +"""DeepSeek-V4-Flash 60k LoRA SFT and vLLM serving on Modal.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from string import Template + +import modal +import modal.experimental + +from dsv4_images import ( + REMOTE_CHECKPOINT_HELPER, + REMOTE_TRAIN_RECIPE, + build_automodel_image, + build_vllm_image, +) + +HERE = Path(__file__).parent + +HF_MODEL = "deepseek-ai/DeepSeek-V4-Flash" +HF_CACHE = "/root/.cache/huggingface" +CHECKPOINTS_DIR = "/checkpoints" + +GPUS_PER_NODE = 8 +N_NODES = int(os.environ.get("N_NODES", "8")) +HOST_MEMORY = (128, 256 * 1024) +SEQ_LENGTH = 60_000 +MAX_STEPS = 5 +CP_SIZE = 16 +PP_SIZE = 4 +MAX_EP_SIZE = 32 +GLOBAL_BATCH_SIZE = 8 +LOCAL_BATCH_SIZE = 4 +LORA_RANK = 64 +LORA_ALPHA = 64 +SYNTHETIC_EXAMPLES = 40 +DIST_BACKEND = "cpu:gloo,cuda:nccl" + +VLLM_ADAPTER_NAME = "deepseek-v4-flash-60k-lora" +VLLM_PORT = 8000 +VLLM_PIPELINE_PARALLEL_SIZE = 4 +VLLM_HF_CONFIG_DIR = Path("/tmp/deepseek-v4-flash-vllm-config") +SERVE_RUN_ID = os.environ.get("SERVE_RUN_ID") +SERVE_MAX_MODEL_LEN = 64 * 1024 + +app = modal.App("example-deepseek-v4-flash-sft") + +hf_cache_vol = modal.Volume.from_name("huggingface-cache", create_if_missing=True) +checkpoints_vol = modal.Volume.from_name( + "example-deepseek-v4-flash-sft-checkpoints", + create_if_missing=True, + version=2, +) + +automodel_image = build_automodel_image(HERE) +vllm_image = build_vllm_image(HERE) + + +def _validate_run_id(run_id: str) -> None: + if not run_id or Path(run_id).name != run_id: + raise ValueError("run_id must be a non-empty path component") + + +def _training_topology(n_nodes: int) -> tuple[int, dict[str, int]]: + if n_nodes < 1: + raise ValueError("N_NODES must be positive") + + total_gpus = n_nodes * GPUS_PER_NODE + if total_gpus % (PP_SIZE * CP_SIZE): + raise ValueError( + f"{total_gpus} GPUs cannot be divided across PP={PP_SIZE} and CP={CP_SIZE}" + ) + + non_pp_size = total_gpus // PP_SIZE + ep_size = min(MAX_EP_SIZE, non_pp_size) + if non_pp_size % ep_size: + raise ValueError(f"{non_pp_size=} must be divisible by {ep_size=}") + + return total_gpus, { + "tp": 1, + "dp": total_gpus // (PP_SIZE * CP_SIZE), + "pp": PP_SIZE, + "cp": CP_SIZE, + "ep": ep_size, + } + + +def _select_uccl_ep_variant() -> str: + """Activate the UCCL extension for the scheduled RDMA provider.""" + import importlib.util + import shutil + + vendors = { + path.read_text().strip() + for path in Path("/sys/class/infiniband").glob("*/device/vendor") + } + if "0x1d0f" in vendors: + provider = "efa" + elif "0x15b3" in vendors: + provider = "mellanox" + else: + raise RuntimeError(f"Unsupported RDMA vendors: {sorted(vendors)}") + + uccl_spec = importlib.util.find_spec("uccl") + if uccl_spec is None or not uccl_spec.submodule_search_locations: + raise RuntimeError("UCCL is missing from the training image") + package_dir = Path(next(iter(uccl_spec.submodule_search_locations))) + source = package_dir / f"ep.{provider}.abi3.so" + if not source.is_file(): + raise RuntimeError(f"Missing UCCL-EP {provider} extension: {source}") + + target = package_dir / "ep.abi3.so" + shutil.copy2(source, target) + print(f"selected_uccl_ep_provider={provider}") + return provider + + +def _write_synthetic_dataset(path: Path) -> None: + """Write fixed-length supervised examples to local storage on each node.""" + path.parent.mkdir(parents=True, exist_ok=True) + payload = " ".join(f"fact_{idx % 4096}" for idx in range(SEQ_LENGTH * 2)) + with path.open("w") as output: + for index in range(SYNTHETIC_EXAMPLES): + row = { + "messages": [ + { + "role": "user", + "content": f"Repeat the synthetic facts for example {index}.", + }, + { + "role": "assistant", + "content": (f"Synthetic facts for example {index}: {payload}"), + }, + ] + } + output.write(json.dumps(row) + "\n") + + +def _render_recipe( + *, + dataset_path: Path, + checkpoint_dir: Path, + ep_size: int, +) -> str: + return Template(Path(REMOTE_TRAIN_RECIPE).read_text()).substitute( + CHECKPOINT_DIR=json.dumps(str(checkpoint_dir)), + CP_SIZE=CP_SIZE, + DATASET_PATH=json.dumps(str(dataset_path)), + DIST_BACKEND=json.dumps(DIST_BACKEND), + EP_SIZE=ep_size, + GLOBAL_BATCH_SIZE=GLOBAL_BATCH_SIZE, + HF_MODEL=json.dumps(HF_MODEL), + LOCAL_BATCH_SIZE=LOCAL_BATCH_SIZE, + LORA_ALPHA=LORA_ALPHA, + LORA_RANK=LORA_RANK, + MAX_STEPS=MAX_STEPS, + PP_SIZE=PP_SIZE, + SEQ_LENGTH=SEQ_LENGTH, + ) + + +@app.function( + image=automodel_image, + gpu="H200:8", + volumes={HF_CACHE: hf_cache_vol, CHECKPOINTS_DIR: checkpoints_vol}, + secrets=[modal.Secret.from_name("huggingface-secret")], + timeout=86400, + retries=0, + memory=HOST_MEMORY, + experimental_options={"efa_enabled": True}, +) +@modal.experimental.clustered(size=N_NODES, rdma=True) +def train_cluster(run_id: str, expected_nodes: int): + import subprocess + + _validate_run_id(run_id) + cluster = modal.experimental.get_cluster_info() + node_rank = cluster.rank + n_nodes = len(cluster.container_ips) if cluster.container_ips else 1 + if n_nodes != expected_nodes: + raise RuntimeError(f"Expected {expected_nodes} nodes, scheduled {n_nodes}") + + total_gpus, topology = _training_topology(n_nodes) + master_addr = cluster.container_ips[0] if cluster.container_ips else "localhost" + os.environ.update( + { + "GLOO_SOCKET_IFNAME": "eth1", + "HF_HOME": HF_CACHE, + "HF_HUB_ENABLE_HF_TRANSFER": "0", + "NCCL_DEBUG": "WARN", + "NCCL_MAX_NCHANNELS": "8", + "NCCL_SOCKET_FAMILY": "AF_INET6", + "NCCL_SOCKET_IFNAME": "eth1", + "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:False", + "UCCL_EP_CPU_TIMEOUT_SECS": "600", + "UCCL_SOCKET_FAMILY": "AF_INET6", + "UCCL_SOCKET_IFNAME": "eth1", + } + ) + + hf_cache_vol.reload() + checkpoints_vol.reload() + uccl_provider = _select_uccl_ep_variant() + + dataset_path = Path("/tmp/dsv4_60k_synthetic/training.jsonl") + _write_synthetic_dataset(dataset_path) + + checkpoint_dir = Path(CHECKPOINTS_DIR) / run_id + if checkpoint_dir.exists(): + raise FileExistsError(f"Checkpoint directory already exists: {checkpoint_dir}") + recipe_path = Path("/tmp/dsv4_train.yaml") + recipe_path.write_text( + _render_recipe( + dataset_path=dataset_path, + checkpoint_dir=checkpoint_dir, + ep_size=topology["ep"], + ) + ) + + summary = { + "run_id": run_id, + "nodes": n_nodes, + "total_gpus": total_gpus, + "topology": topology, + "sequence_length": SEQ_LENGTH, + "optimizer_steps": MAX_STEPS, + "checkpoint_dir": str(checkpoint_dir), + "uccl_provider": uccl_provider, + } + if node_rank == 0: + print(json.dumps(summary, indent=2, sort_keys=True)) + + command = [ + "torchrun", + "--nproc-per-node", + str(GPUS_PER_NODE), + "--nnodes", + str(n_nodes), + "--node-rank", + str(node_rank), + "--master-addr", + master_addr, + "--master-port", + "29501", + "-m", + "nemo_automodel.cli.app", + str(recipe_path), + ] + subprocess.run(command, check=True) + checkpoints_vol.commit() + return summary + + +@app.function( + image=automodel_image, + volumes={HF_CACHE: hf_cache_vol, CHECKPOINTS_DIR: checkpoints_vol}, + secrets=[modal.Secret.from_name("huggingface-secret")], + timeout=1800, + memory=HOST_MEMORY, +) +def finalize(run_id: str): + """Merge the four pipeline-stage shards into a serving adapter.""" + import sys + + sys.path.insert(0, str(Path(REMOTE_CHECKPOINT_HELPER).parent)) + from dsv4_checkpoint import finalize_adapter + + checkpoints_vol.reload() + manifest = finalize_adapter(CHECKPOINTS_DIR, run_id) + checkpoints_vol.commit() + print(json.dumps(manifest, indent=2, sort_keys=True)) + return manifest + + +def _download_vllm_config() -> None: + from huggingface_hub import hf_hub_download + + config_path = Path( + hf_hub_download( + repo_id=HF_MODEL, + filename="config.json", + cache_dir="/tmp/vllm-hf-config-cache", + local_dir=VLLM_HF_CONFIG_DIR, + force_download=True, + token=os.environ.get("HF_TOKEN"), + ) + ) + config = json.loads(config_path.read_text()) + expected_quantization = { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "fp8", + "scale_fmt": "ue8m0", + "weight_block_size": [128, 128], + } + if config.get("quantization_config") != expected_quantization: + raise RuntimeError( + "Unexpected DeepSeek-V4-Flash quantization configuration: " + f"{config.get('quantization_config')}" + ) + + +def _vllm_command(adapter_dir: Path) -> list[str]: + return [ + "python", + "-m", + "vllm.entrypoints.openai.api_server", + "--model", + HF_MODEL, + "--served-model-name", + HF_MODEL, + "--pipeline-parallel-size", + str(VLLM_PIPELINE_PARALLEL_SIZE), + "--enable-expert-parallel", + "--trust-remote-code", + "--hf-config-path", + str(VLLM_HF_CONFIG_DIR), + "--tokenizer-mode", + "deepseek_v4", + "--reasoning-parser", + "deepseek_v4", + "--max-model-len", + str(SERVE_MAX_MODEL_LEN), + "--max-num-seqs", + "1", + "--max-num-batched-tokens", + "8192", + "--gpu-memory-utilization", + "0.92", + "--kv-cache-dtype", + "fp8", + "--block-size", + "256", + "--enable-lora", + "--max-loras", + "1", + "--max-lora-rank", + str(LORA_RANK), + "--lora-target-modules", + "fused_wqa_wkv", + "wq_b", + "--lora-modules", + f"{VLLM_ADAPTER_NAME}={adapter_dir}", + "--enforce-eager", + "--no-enable-flashinfer-autotune", + "--no-enable-log-requests", + "--disable-uvicorn-access-log", + "--port", + str(VLLM_PORT), + ] + + +@app.function( + image=vllm_image, + gpu=f"H200:{VLLM_PIPELINE_PARALLEL_SIZE}", + volumes={HF_CACHE: hf_cache_vol, CHECKPOINTS_DIR: checkpoints_vol}, + secrets=[modal.Secret.from_name("huggingface-secret")], + timeout=86400, + memory=HOST_MEMORY, + experimental_options={"efa_enabled": True}, +) +@modal.web_server( + VLLM_PORT, + startup_timeout=3600, + requires_proxy_auth=True, +) +def serve(): + """Serve the finalized adapter selected by SERVE_RUN_ID.""" + import subprocess + import sys + + if not SERVE_RUN_ID: + raise RuntimeError("Set SERVE_RUN_ID to a finalized checkpoint run ID") + + sys.path.insert(0, str(Path(REMOTE_CHECKPOINT_HELPER).parent)) + from dsv4_checkpoint import finalized_adapter + + hf_cache_vol.reload() + checkpoints_vol.reload() + _download_vllm_config() + adapter_dir, manifest = finalized_adapter(CHECKPOINTS_DIR, SERVE_RUN_ID) + print(f"adapter_sha256={manifest['adapter_sha256']}") + subprocess.Popen(_vllm_command(adapter_dir)) + + +@app.local_entrypoint() +def train(run_id: str): + """Start a detached multi-node training call.""" + _validate_run_id(run_id) + _, topology = _training_topology(N_NODES) + call = train_cluster.spawn(run_id=run_id, expected_nodes=N_NODES) + call_id = ( + getattr(call, "object_id", None) + or getattr(call, "function_call_id", None) + or str(call) + ) + print( + json.dumps( + { + "function_call_id": call_id, + "run_id": run_id, + "nodes": N_NODES, + "topology": topology, + }, + indent=2, + ) + ) diff --git a/deepseek-v4-flash-sft/patches/automodel_checkpoint_dequant.patch b/deepseek-v4-flash-sft/patches/automodel_checkpoint_dequant.patch new file mode 100644 index 0000000..84eb139 --- /dev/null +++ b/deepseek-v4-flash-sft/patches/automodel_checkpoint_dequant.patch @@ -0,0 +1,86 @@ +diff --git a/nemo_automodel/_transformers/auto_model.py b/nemo_automodel/_transformers/auto_model.py +index cbc6bad617..6e9e6b9932 100644 +--- a/nemo_automodel/_transformers/auto_model.py ++++ b/nemo_automodel/_transformers/auto_model.py +@@ -416,6 +416,7 @@ class _BaseNeMoAutoModelClass: + has_packed_sequence = kwargs.pop("has_packed_sequence", False) + freeze_config = kwargs.pop("freeze_config", None) + cache_dir = kwargs.pop("cache_dir", hf_constants.HF_HUB_CACHE) ++ dequantize_base_checkpoint = kwargs.pop("dequantize_base_checkpoint", None) + + def _retry(**override): + """Re-enter ``_build_model`` with overridden parameters.""" +@@ -426,6 +427,7 @@ class _BaseNeMoAutoModelClass: + "has_packed_sequence": has_packed_sequence, + "freeze_config": freeze_config, + "cache_dir": cache_dir, ++ "dequantize_base_checkpoint": dequantize_base_checkpoint, + } + return cls._build_model( + pretrained_model_name_or_path_or_config, +@@ -629,6 +631,7 @@ class _BaseNeMoAutoModelClass: + compile_config=compile_config, + load_base_model=load_base_model, + cache_dir=cache_dir, ++ dequantize_base_checkpoint=dequantize_base_checkpoint, + freeze_config=freeze_config, + weights_already_loaded=weights_already_loaded, + inject_te_attention=inject_te_attention, +diff --git a/nemo_automodel/_transformers/infrastructure.py b/nemo_automodel/_transformers/infrastructure.py +index 8f2d1f693d..889d197b75 100644 +--- a/nemo_automodel/_transformers/infrastructure.py ++++ b/nemo_automodel/_transformers/infrastructure.py +@@ -395,6 +395,7 @@ def apply_model_infrastructure( + parallelize_fn=None, + compile_config=None, + load_base_model=False, ++ dequantize_base_checkpoint: bool | None = None, + cache_dir=None, + pretrained_model_name_or_path="", + weights_already_loaded=False, +@@ -427,6 +428,8 @@ def apply_model_infrastructure( + compile_config: Compilation configuration. Default: None + pretrained_model_name_or_path: Model name or path for checkpoint loading. Default: "" + load_base_model: Whether to load base model weights (True for from_pretrained). Default: False ++ dequantize_base_checkpoint: Whether the state-dict adapter should construct ++ quantized checkpoint placeholders before dequantizing loaded weights. + cache_dir: Cache directory for model weights. Default: None + weights_already_loaded: Whether pretrained weights were already loaded during + model init (e.g., by HF's from_pretrained on a real device, which also +@@ -455,6 +458,7 @@ def apply_model_infrastructure( + model_repo_id=pretrained_model_name_or_path, + save_consolidated=False, + is_peft=peft_config is not None, ++ dequantize_base_checkpoint=dequantize_base_checkpoint, + ) + checkpointer = Checkpointer( + ckpt_config, +diff --git a/nemo_automodel/recipes/llm/train_ft.py b/nemo_automodel/recipes/llm/train_ft.py +index fa3cdba275..f89ed92b45 100644 +--- a/nemo_automodel/recipes/llm/train_ft.py ++++ b/nemo_automodel/recipes/llm/train_ft.py +@@ -207,6 +207,7 @@ def build_model( + unfreeze_modules: list[str] | None = None, + sdpa_method: list[str] | None = None, + device_mesh=None, ++ dequantize_base_checkpoint: bool | None = None, + ) -> tuple[nn.Module | AutoPipeline, list["Optimizer"]]: # noqa: F821 + """Build and initialize a model. + +@@ -236,6 +237,8 @@ def build_model( + kwargs["distributed_setup"] = distributed_setup + elif device_mesh is not None: + kwargs["device_mesh"] = device_mesh ++ if dequantize_base_checkpoint is not None: ++ kwargs["dequantize_base_checkpoint"] = dequantize_base_checkpoint + + if cfg_qat is not None and cfg_qat.get("enabled", False): + if cfg_peft is not None: +@@ -593,6 +596,7 @@ class TrainFinetuneRecipeForNextTokenPrediction(BaseRecipe): + distributed_setup=self.distributed_setup, + cfg_qat=self.cfg.get("qat", None), + sdpa_method=self.cfg.get("sdpa_method", None), ++ dequantize_base_checkpoint=checkpoint_config.dequantize_base_checkpoint, + ) + optimizer = self.cfg.optimizer.build(model, device_mesh=self.device_mesh, is_peft=self.peft_config is not None) + allow_megatron_fsdp_sharding = getattr(self.cfg.optimizer, "supports_megatron_fsdp_sharding", True) diff --git a/deepseek-v4-flash-sft/patches/automodel_composite_backend.patch b/deepseek-v4-flash-sft/patches/automodel_composite_backend.patch new file mode 100644 index 0000000..85d1d40 --- /dev/null +++ b/deepseek-v4-flash-sft/patches/automodel_composite_backend.patch @@ -0,0 +1,22 @@ +diff --git a/nemo_automodel/components/training/signal_handler.py b/nemo_automodel/components/training/signal_handler.py +index e1a526baad..a7f715e765 100644 +--- a/nemo_automodel/components/training/signal_handler.py ++++ b/nemo_automodel/components/training/signal_handler.py +@@ -38,11 +38,13 @@ def get_device(local_rank: Optional[int] = None) -> torch.device: +- backend = torch.distributed.get_backend() +- if backend == "nccl": ++ backend = str(torch.distributed.get_backend()).lower() ++ if backend == "nccl": + if local_rank is None: + device = torch.device("cuda") + else: + device = torch.device(f"cuda:{local_rank}") +- elif backend == "gloo": ++ elif backend == "gloo" or "gloo" in backend or "undefined" in backend: ++ # Composite process groups dispatch collectives by tensor device, so ++ # keep this control-plane collective on CPU/Gloo. + device = torch.device("cpu") + else: +- raise RuntimeError ++ raise RuntimeError(f"Unsupported distributed backend: {backend!r}") + return device diff --git a/deepseek-v4-flash-sft/patches/automodel_pp_peft_checkpoint.patch b/deepseek-v4-flash-sft/patches/automodel_pp_peft_checkpoint.patch new file mode 100644 index 0000000..b2a2bdc --- /dev/null +++ b/deepseek-v4-flash-sft/patches/automodel_pp_peft_checkpoint.patch @@ -0,0 +1,79 @@ +diff --git a/nemo_automodel/components/checkpoint/checkpointing.py b/nemo_automodel/components/checkpoint/checkpointing.py +--- a/nemo_automodel/components/checkpoint/checkpointing.py ++++ b/nemo_automodel/components/checkpoint/checkpointing.py +@@ -1139,7 +1139,8 @@ class Checkpointer: + """ + Save a state dictionary to `path` using DCP or PEFT special-case logic. + +- - For PEFT model saves: only rank 0 writes `adapter_model.safetensors`. ++ - For PEFT model saves without PP: rank 0 writes `adapter_model.safetensors`. ++ - For PEFT model saves with PP: one rank per stage writes a stage shard. + - If async mode is enabled, schedule an asynchronous save. + + Args: +@@ -1156,10 +1157,37 @@ class Checkpointer: + is_model = True if "/model" in path else False + # PEFT saving is done on rank0 so it is a special case + if self.config.is_peft and is_model: +- if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0: ++ if not torch.distributed.is_initialized(): + _save_safetensors(state_dict, _adapter_path(path)) +- if torch.distributed.is_initialized(): +- torch.distributed.barrier() ++ return ++ ++ rank = torch.distributed.get_rank() ++ world_size = torch.distributed.get_world_size() ++ rank_and_stage = [None] * world_size ++ torch.distributed.all_gather_object( ++ rank_and_stage, (rank, self.pp_rank) ++ ) ++ writer_by_stage = {} ++ for global_rank, pp_rank in rank_and_stage: ++ writer_by_stage.setdefault(pp_rank, global_rank) ++ writer_by_stage[pp_rank] = min( ++ writer_by_stage[pp_rank], global_rank ++ ) ++ ++ if len(writer_by_stage) == 1: ++ if rank == 0: ++ _save_safetensors(state_dict, _adapter_path(path)) ++ elif rank == writer_by_stage[self.pp_rank]: ++ adapter_path = _adapter_path(path) ++ stem, suffix = os.path.splitext(adapter_path) ++ stage_path = ( ++ f"{stem}.pp-{self.pp_rank:05d}-of-" ++ f"{len(writer_by_stage):05d}{suffix}" ++ ) ++ _save_safetensors(state_dict, stage_path) ++ logger.info("Saved pipeline-stage PEFT adapter to %s", stage_path) ++ ++ torch.distributed.barrier() + return + + ret = None +diff --git a/nemo_automodel/components/checkpoint/config.py b/nemo_automodel/components/checkpoint/config.py +--- a/nemo_automodel/components/checkpoint/config.py ++++ b/nemo_automodel/components/checkpoint/config.py +@@ -103,6 +103,7 @@ class CheckpointingConfig: + is_peft: bool = False + model_state_dict_keys: list[str] | None = ( + None # copy of the model state dict keys before any parallelization. Kept for BW compatibility. + ) ++ save_optimizer: bool = True + is_async: bool = False + dequantize_base_checkpoint: bool | None = None + original_model_root_dir: str | None = None +diff --git a/nemo_automodel/recipes/base_recipe.py b/nemo_automodel/recipes/base_recipe.py +--- a/nemo_automodel/recipes/base_recipe.py ++++ b/nemo_automodel/recipes/base_recipe.py +@@ -445,7 +445,8 @@ class BaseRecipe: + for opt in optimizers: + if hasattr(opt, "synchronize_for_checkpoint"): + opt.synchronize_for_checkpoint() +- self.checkpointer.save_optimizer(optimizer, model, path, scheduler) ++ if self.checkpointer.config.save_optimizer: ++ self.checkpointer.save_optimizer(optimizer, model, path, scheduler) + save_config(config.raw_config, path) + if is_dist_initialized: + torch.distributed.barrier() diff --git a/deepseek-v4-flash-sft/patches/automodel_uccl_teardown.patch b/deepseek-v4-flash-sft/patches/automodel_uccl_teardown.patch new file mode 100644 index 0000000..3ef78b2 --- /dev/null +++ b/deepseek-v4-flash-sft/patches/automodel_uccl_teardown.patch @@ -0,0 +1,56 @@ +diff --git a/nemo_automodel/components/distributed/init_utils.py b/nemo_automodel/components/distributed/init_utils.py +index 57433706c0..98852b1dbd 100644 +--- a/nemo_automodel/components/distributed/init_utils.py ++++ b/nemo_automodel/components/distributed/init_utils.py +@@ -157,19 +157,18 @@ def destroy_global_state(): + It temporarily ignores SIGINT to avoid interruption during cleanup. + +- For MoE runs that use DeepEP, the process-global DeepEP ``Buffer`` (NVSHMEM symmetric memory +- plus its own NCCL sub-groups) is destroyed *before* ``destroy_process_group()``. Tearing the +- buffer down first releases that pending collective state; otherwise ``destroy_process_group()`` +- hangs on it. Freeing the buffer is a no-op when DeepEP was never used. ++ For MoE runs, process-global DeepEP or UCCL-EP buffers are destroyed *before* ++ ``destroy_process_group()`` so their communication runtimes do not outlive CUDA. + """ + # Don't allow Ctrl+C to interrupt this handler + signal.signal(signal.SIGINT, signal.SIG_IGN) + try: +- from nemo_automodel.components.moe.megatron.fused_a2a import free_buffer ++ from nemo_automodel.components.moe.megatron.fused_a2a import free_buffer, free_uccl_buffer + + if torch.cuda.is_available(): + torch.cuda.synchronize() + free_buffer() +- except Exception: # pragma: no cover - best effort; deep_ep may be absent ++ free_uccl_buffer() ++ except Exception: # pragma: no cover - best effort; communication libraries may be absent + pass + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.destroy_process_group() +diff --git a/nemo_automodel/components/moe/megatron/fused_a2a.py b/nemo_automodel/components/moe/megatron/fused_a2a.py +index cec285cc1e..a70a59ab9c 100644 +--- a/nemo_automodel/components/moe/megatron/fused_a2a.py ++++ b/nemo_automodel/components/moe/megatron/fused_a2a.py +@@ -541,10 +541,21 @@ def get_uccl_buffer(group: torch.distributed.ProcessGroup, hidden_bytes: int): + or _uccl_buffer.num_nvl_bytes < num_nvl_bytes + or _uccl_buffer.num_rdma_bytes < num_rdma_bytes + ): +- _uccl_buffer = UCCLBuffer(group, num_nvl_bytes, num_rdma_bytes) ++ _uccl_buffer = UCCLBuffer( ++ group, num_nvl_bytes, num_rdma_bytes, explicitly_destroy=True ++ ) + return _uccl_buffer + + ++def free_uccl_buffer() -> None: ++ """Destroy the global UCCL-EP buffer while CUDA is still available.""" ++ global _uccl_buffer ++ buffer = _uccl_buffer ++ _uccl_buffer = None ++ if buffer is not None: ++ buffer.destroy() ++ ++ + class UCCLFusedDispatch(torch.autograd.Function): + """Fused dispatch using UCCL-EP instead of DeepEP.""" + diff --git a/deepseek-v4-flash-sft/patches/uccl_ipv6_oob.patch b/deepseek-v4-flash-sft/patches/uccl_ipv6_oob.patch new file mode 100644 index 0000000..f5e9081 --- /dev/null +++ b/deepseek-v4-flash-sft/patches/uccl_ipv6_oob.patch @@ -0,0 +1,259 @@ +--- a/ep/src/rdma.cpp ++++ b/ep/src/rdma.cpp +@@ -318,12 +318,12 @@ + int listen_fd, + RDMAConnectionInfo* remote_array) { + int sockfd; +- struct sockaddr_in addr; ++ struct sockaddr_storage addr; + memset(&addr, 0, sizeof(addr)); + + // Already listening when calling uccl::create_listen_socket(). + socklen_t len = sizeof(addr); +- sockfd = accept(listen_fd, (struct sockaddr*)&addr, &len); ++ sockfd = accept(listen_fd, reinterpret_cast(&addr), &len); + + // Exchange info + uccl::receive_message(sockfd, actual_peer, sizeof(*actual_peer)); +@@ -335,20 +335,18 @@ + void send_connection_info_as_client(int my_rank, int peer, char const* peer_ip, + int peer_listen_port, + RDMAConnectionInfo* local) { +- int sockfd; +- struct sockaddr_in addr; +- memset(&addr, 0, sizeof(addr)); ++ struct sockaddr_storage addr; ++ socklen_t const addr_len = ++ uccl::make_oob_socket_address(peer_ip, peer_listen_port, &addr); + + // Connect +- sockfd = socket(AF_INET, SOCK_STREAM, 0); ++ int sockfd = socket(addr.ss_family, SOCK_STREAM, 0); + int one = 1; + setsockopt(sockfd, SOL_SOCKET, SO_REUSEADDR, &one, sizeof(one)); +- addr.sin_family = AF_INET; +- addr.sin_port = htons(peer_listen_port); +- inet_pton(AF_INET, peer_ip, &addr.sin_addr); + + int retry = 0; +- while (connect(sockfd, (struct sockaddr*)&addr, sizeof(addr)) != 0) { ++ while (connect(sockfd, reinterpret_cast(&addr), addr_len) != ++ 0) { + if (errno == ECONNREFUSED || errno == ENETUNREACH) { + if (++retry > MAX_RETRIES) { + fprintf(stderr, "Rank %d: failed to connect to %d after %d retries\n", +--- a/include/util/net.h ++++ b/include/util/net.h +@@ -130,7 +130,8 @@ + /* Allow the user to force the IPv4/IPv6 interface selection */ + static inline int env_socket_family(void) { + int family = -1; // Family selection is not forced, will use first one found +- char* env = getenv("NCCL_SOCKET_FAMILY"); ++ char* env = getenv("UCCL_SOCKET_FAMILY"); ++ if (env == NULL) env = getenv("NCCL_SOCKET_FAMILY"); + if (env == NULL) return family; + + UCCL_LOG(INFO, UCCL_UTIL) +@@ -497,6 +498,16 @@ + int num_ifs = + uccl::find_interfaces(uccl_ifname, &uccl_ifaddr, MAX_IF_NAME_SIZE, 1); + UCCL_CHECK(num_ifs == 1) << "No IP interface found"; +- return uccl::get_dev_ip(uccl_ifname); ++ ++ int const family = uccl_ifaddr.sa.sa_family; ++ void const* addr = family == AF_INET ++ ? static_cast(&uccl_ifaddr.sin.sin_addr) ++ : static_cast(&uccl_ifaddr.sin6.sin6_addr); ++ char address[INET6_ADDRSTRLEN]{}; ++ UCCL_CHECK(family == AF_INET || family == AF_INET6) ++ << "Unsupported OOB address family " << family; ++ UCCL_CHECK(inet_ntop(family, addr, address, sizeof(address)) != nullptr) ++ << "Failed to format OOB address for interface " << uccl_ifname; ++ return std::string(address); + } +-} // namespace uccl +\ No newline at end of file ++} // namespace uccl +--- a/include/util/util.h ++++ b/include/util/util.h +@@ -169,20 +169,73 @@ + UCCL_DCHECK(ret == sizeof(bool) && sync) << ret << ", " << sync; + } + ++inline int oob_socket_family() { ++ char const* family = getenv("UCCL_SOCKET_FAMILY"); ++ if (family == nullptr) family = getenv("NCCL_SOCKET_FAMILY"); ++ return family != nullptr && strcmp(family, "AF_INET6") == 0 ? AF_INET6 ++ : AF_INET; ++} ++ ++inline socklen_t make_oob_any_address(int family, uint16_t port, ++ struct sockaddr_storage* storage) { ++ memset(storage, 0, sizeof(*storage)); ++ if (family == AF_INET6) { ++ auto* addr = reinterpret_cast(storage); ++ addr->sin6_family = AF_INET6; ++ addr->sin6_addr = in6addr_any; ++ addr->sin6_port = htons(port); ++ return sizeof(*addr); ++ } ++ ++ auto* addr = reinterpret_cast(storage); ++ addr->sin_family = AF_INET; ++ addr->sin_addr.s_addr = INADDR_ANY; ++ addr->sin_port = htons(port); ++ return sizeof(*addr); ++} ++ ++inline socklen_t make_oob_socket_address(char const* ip, uint16_t port, ++ struct sockaddr_storage* storage) { ++ memset(storage, 0, sizeof(*storage)); ++ if (strchr(ip, ':') != nullptr) { ++ auto* addr = reinterpret_cast(storage); ++ addr->sin6_family = AF_INET6; ++ addr->sin6_port = htons(port); ++ UCCL_CHECK(inet_pton(AF_INET6, ip, &addr->sin6_addr) == 1) ++ << "Invalid IPv6 OOB address " << ip; ++ return sizeof(*addr); ++ } ++ ++ auto* addr = reinterpret_cast(storage); ++ addr->sin_family = AF_INET; ++ addr->sin_port = htons(port); ++ UCCL_CHECK(inet_pton(AF_INET, ip, &addr->sin_addr) == 1) ++ << "Invalid IPv4 OOB address " << ip; ++ return sizeof(*addr); ++} ++ ++inline uint16_t oob_socket_port(struct sockaddr_storage const& storage) { ++ if (storage.ss_family == AF_INET6) { ++ auto const* addr = reinterpret_cast(&storage); ++ return ntohs(addr->sin6_port); ++ } ++ auto const* addr = reinterpret_cast(&storage); ++ return ntohs(addr->sin_port); ++} ++ + inline void create_listen_socket(int* listen_fd, uint16_t listen_port) { +- *listen_fd = socket(AF_INET, SOCK_STREAM, 0); ++ int const family = oob_socket_family(); ++ *listen_fd = socket(family, SOCK_STREAM, 0); + UCCL_DCHECK(*listen_fd >= 0) << "ERROR: opening socket"; + int flag = 1; + UCCL_CHECK( + setsockopt(*listen_fd, SOL_SOCKET, SO_REUSEADDR, &flag, sizeof(int)) >= 0) + << "ERROR: setsockopt SO_REUSEADDR fails"; +- struct sockaddr_in serv_addr; +- bzero((char*)&serv_addr, sizeof(serv_addr)); +- serv_addr.sin_family = AF_INET; +- serv_addr.sin_addr.s_addr = INADDR_ANY; +- serv_addr.sin_port = htons(listen_port); +- UCCL_CHECK( +- bind(*listen_fd, (struct sockaddr*)&serv_addr, sizeof(serv_addr)) >= 0) ++ struct sockaddr_storage serv_addr; ++ socklen_t const addr_len = ++ make_oob_any_address(family, listen_port, &serv_addr); ++ UCCL_CHECK(bind(*listen_fd, reinterpret_cast(&serv_addr), ++ addr_len) >= 0) + << "ERROR: binding"; + + UCCL_CHECK(!listen(*listen_fd, 128)) << "ERROR: listen"; +@@ -190,7 +243,8 @@ + } + + inline uint16_t create_listen_socket(int* listen_fd) { +- *listen_fd = socket(AF_INET, SOCK_STREAM, 0); ++ int const family = oob_socket_family(); ++ *listen_fd = socket(family, SOCK_STREAM, 0); + UCCL_DCHECK(*listen_fd >= 0) << "ERROR: opening socket"; + + int flag = 1; +@@ -198,22 +252,20 @@ + setsockopt(*listen_fd, SOL_SOCKET, SO_REUSEADDR, &flag, sizeof(int)) >= 0) + << "ERROR: setsockopt SO_REUSEADDR fails"; + +- struct sockaddr_in serv_addr; +- bzero((char*)&serv_addr, sizeof(serv_addr)); +- serv_addr.sin_family = AF_INET; +- serv_addr.sin_addr.s_addr = INADDR_ANY; +- serv_addr.sin_port = htons(0); // Ask OS for ephemeral port +- +- UCCL_CHECK( +- bind(*listen_fd, (struct sockaddr*)&serv_addr, sizeof(serv_addr)) >= 0) ++ struct sockaddr_storage serv_addr; ++ socklen_t addr_len = make_oob_any_address(family, 0, &serv_addr); ++ UCCL_CHECK(bind(*listen_fd, reinterpret_cast(&serv_addr), ++ addr_len) >= 0) + << "ERROR: binding"; + + // Get the assigned port +- socklen_t len = sizeof(serv_addr); +- UCCL_CHECK(getsockname(*listen_fd, (struct sockaddr*)&serv_addr, &len) >= 0) ++ addr_len = sizeof(serv_addr); ++ UCCL_CHECK(getsockname(*listen_fd, ++ reinterpret_cast(&serv_addr), ++ &addr_len) >= 0) + << "ERROR: getsockname"; + +- uint16_t assigned_port = ntohs(serv_addr.sin_port); ++ uint16_t assigned_port = oob_socket_port(serv_addr); + + UCCL_CHECK(!listen(*listen_fd, 128)) << "ERROR: listen"; + UCCL_VLOG(5) << "[Endpoint] server ready, listening on ephemeral port " +@@ -231,19 +283,29 @@ + UCCL_VLOG(5) << "[listen_accept_exchange] server ready, listening on port " + << oobport; + +- struct sockaddr_in client_addr; ++ struct sockaddr_storage client_addr; + socklen_t client_len = sizeof(client_addr); + + int client_fd = +- accept(listen_fd, (struct sockaddr*)&client_addr, &client_len); ++ accept(listen_fd, reinterpret_cast(&client_addr), ++ &client_len); + UCCL_CHECK(client_fd >= 0) << "Failed to accept connection"; + + // Set nodelay + int flag = 1; + setsockopt(client_fd, IPPROTO_TCP, TCP_NODELAY, (void*)&flag, sizeof(int)); + +- char client_ip[INET_ADDRSTRLEN]; +- inet_ntop(AF_INET, &(client_addr.sin_addr), client_ip, INET_ADDRSTRLEN); ++ char client_ip[INET6_ADDRSTRLEN]{}; ++ void const* client_addr_ptr = ++ client_addr.ss_family == AF_INET6 ++ ? static_cast(&reinterpret_cast( ++ &client_addr) ++ ->sin6_addr) ++ : static_cast(&reinterpret_cast( ++ &client_addr) ++ ->sin_addr); ++ inet_ntop(client_addr.ss_family, client_addr_ptr, client_ip, ++ sizeof(client_ip)); + UCCL_VLOG(5) << "[listen_accept_exchange] accepted connection from " + << client_ip; + +@@ -257,16 +319,14 @@ + inline static void connect_exchange(int oobport, std::string oob_ip, + void* send_data, int send_size, + void* recv_data, int recv_size) { +- int sockfd = socket(AF_INET, SOCK_STREAM, 0); ++ struct sockaddr_storage server_addr; ++ socklen_t const addr_len = ++ make_oob_socket_address(oob_ip.c_str(), oobport, &server_addr); ++ int sockfd = socket(server_addr.ss_family, SOCK_STREAM, 0); + UCCL_CHECK(sockfd >= 0) << "Failed to create socket"; + +- struct sockaddr_in server_addr; +- server_addr.sin_family = AF_INET; +- server_addr.sin_port = htons(oobport); +- server_addr.sin_addr.s_addr = inet_addr(oob_ip.c_str()); +- +- while (connect(sockfd, (struct sockaddr*)&server_addr, sizeof(server_addr)) < +- 0) { ++ while (connect(sockfd, reinterpret_cast(&server_addr), ++ addr_len) < 0) { + std::this_thread::sleep_for(std::chrono::milliseconds(1000)); + UCCL_VLOG(5) << "[connect_exchange] connecting to " << oob_ip << ":" + << oobport; diff --git a/deepseek-v4-flash-sft/patches/vllm_deepseek_v4_lora.patch b/deepseek-v4-flash-sft/patches/vllm_deepseek_v4_lora.patch new file mode 100644 index 0000000..73133c2 --- /dev/null +++ b/deepseek-v4-flash-sft/patches/vllm_deepseek_v4_lora.patch @@ -0,0 +1,46 @@ +diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py +index 5f4a638c24..9d11e35f53 100644 +--- a/vllm/models/deepseek_v4/nvidia/model.py ++++ b/vllm/models/deepseek_v4/nvidia/model.py +@@ -52,6 +52,7 @@ from vllm.model_executor.models.interfaces import ( + EagleModelMixin, + MixtureOfExperts, + SupportsEagle3, ++ SupportsLoRA, + SupportsPP, + ) + from vllm.model_executor.models.utils import ( +@@ -1312,6 +1313,8 @@ def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: + }, + orig_to_new_substr={ + ".shared_experts.w2": ".shared_experts.down_proj", ++ ".self_attn.": ".attn.", ++ ".q_b_proj": ".wq_b", + }, + ) + +@@ -1354,8 +1357,15 @@ class DeepseekV4MixtureOfExperts(MixtureOfExperts): + + + class DeepseekV4ForCausalLM( +- nn.Module, SupportsPP, SupportsEagle3, DeepseekV4MixtureOfExperts ++ nn.Module, ++ SupportsPP, ++ SupportsLoRA, ++ SupportsEagle3, ++ DeepseekV4MixtureOfExperts, + ): ++ lora_packed_modules_mapping = { ++ "fused_wqa_wkv": ["q_a_proj", "kv_proj"], ++ } + model_cls = DeepseekV4Model + + # Default mapper assumes the original FP4-expert checkpoint layout. +@@ -1364,6 +1374,7 @@ class DeepseekV4ForCausalLM( + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() ++ self.packed_modules_mapping = self.lora_packed_modules_mapping + + config = vllm_config.model_config.hf_config + self.config = config diff --git a/deepseek-v4-flash-sft/train_recipe.yaml b/deepseek-v4-flash-sft/train_recipe.yaml new file mode 100644 index 0000000..95d6064 --- /dev/null +++ b/deepseek-v4-flash-sft/train_recipe.yaml @@ -0,0 +1,119 @@ +recipe: TrainFinetuneRecipeForNextTokenPrediction + +seed: 1234 + +step_scheduler: + global_batch_size: $GLOBAL_BATCH_SIZE + local_batch_size: $LOCAL_BATCH_SIZE + ckpt_every_steps: $MAX_STEPS + save_checkpoint_every_epoch: false + val_every_steps: 100000 + gc_every_steps: 1 + num_epochs: 1 + max_steps: $MAX_STEPS + +distributed: + strategy: fsdp2 + tp_size: 1 + cp_size: $CP_SIZE + pp_size: $PP_SIZE + ep_size: $EP_SIZE + sequence_parallel: false + activation_checkpointing: true + pipeline: + pp_schedule: 1f1b + pp_microbatch_size: 1 + layers_per_stage: 11 + round_virtual_stages_to_pp_multiple: down + scale_grads_in_schedule: false + patch_inner_model: false + patch_causal_lm_model: false + moe: + ignore_router_for_ac: true + reshard_after_forward: false + wrap_outer_model: false + +dist_env: + backend: $DIST_BACKEND + timeout_minutes: 60 + +model: + _target_: nemo_automodel.NeMoAutoModelForCausalLM.from_config + config: + _target_: nemo_automodel.components.models.deepseek_v4.config.DeepseekV4Config.from_pretrained + pretrained_model_name_or_path: $HF_MODEL + name_or_path: $HF_MODEL + num_nextn_predict_layers: 0 + trust_remote_code: false + load_base_model: true + mtp_loss_scaling_factor: 0.1 + backend: + _target_: nemo_automodel.components.models.common.BackendConfig + attn: tilelang + linear: torch + rms_norm: torch_fp32 + rope_fusion: false + dispatcher: uccl_ep + experts: torch_mm + enable_hf_state_dict_adapter: true + enable_fsdp_optimizations: true + +peft: + _target_: nemo_automodel.components._peft.lora.PeftConfig + target_modules: + - "*.self_attn.wq_a" + - "*.self_attn.wq_b" + - "*.self_attn.wkv" + dim: $LORA_RANK + alpha: $LORA_ALPHA + dropout: 0.0 + use_memory_efficient_lora: true + use_triton: false + +checkpoint: + enabled: true + checkpoint_dir: $CHECKPOINT_DIR + model_save_format: safetensors + save_consolidated: final + save_optimizer: false + dequantize_base_checkpoint: true + +loss_fn: + _target_: nemo_automodel.components.loss.masked_ce.MaskedCrossEntropy + +dataset: + _target_: nemo_automodel.components.datasets.llm.chat_dataset.ChatDataset + path_or_dataset_id: $DATASET_PATH + split: train + shuffle_seed: 42 + truncation: true + seq_length: $SEQ_LENGTH + padding: max_length + chat_template: |- + {% for message in messages %}{% if message['role'] == 'user' %}User: + {{ message['content'] }} + {% elif message['role'] == 'assistant' %}Assistant: + {% generation %}{{ message['content'] }}{% endgeneration %} + {% endif %}{% endfor %} + tokenizer: + _target_: transformers.AutoTokenizer.from_pretrained + pretrained_model_name_or_path: $HF_MODEL + +packed_sequence: + packed_sequence_size: 0 + +dataloader: + _target_: torchdata.stateful_dataloader.StatefulDataLoader + collate_fn: + _target_: nemo_automodel.components.datasets.utils.default_collater + pad_seq_len_divisible: 8 + shuffle: true + +optimizer: + _target_: torch.optim.AdamW + betas: + - 0.9 + - 0.95 + eps: 1e-8 + lr: 1e-5 + weight_decay: 0.1