diff --git a/miles/configs/glm5_2_744b_a40b_lora.py b/miles/configs/glm5_2_744b_a40b_lora.py new file mode 100644 index 0000000..db5be4c --- /dev/null +++ b/miles/configs/glm5_2_744b_a40b_lora.py @@ -0,0 +1,190 @@ +"""GLM-5.2 (full 744B-A40B) LoRA GRPO — 8 nodes x 8 H200, colocated. + +this configuration is adapted from: +https://github.com/radixark/miles/blob/c4d9d49cbf8a39185f4c80c0f6084836fc759819/launch_glm_rl_att_unfused_moe.sh + +The launcher was later deleted. This config uses its Megatron/BSHD branch with +a smaller padding quantum to control activation memory. The topology is 8 nodes +x 8 H200, EP 32, DP 8, TP 8, PP 1, CP 1. + +on GSm8k dataset with max response length 256. + +to run: + + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora uv run modal run miles/modal_train.py::download_model + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora uv run modal run miles/modal_train.py::download_data + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora uv run modal run miles/modal_train.py::train +""" + +from configs.base import ModalConfig, MilesConfig, DATA_PATH, CHECKPOINTS_PATH, HF_CACHE_PATH + +modal = ModalConfig( + docker_image="radixark/miles:dev-202607090055", # validated versions sglang 0.5.15, Megatron-Bridge 0.5.0, PR #1559 + #1593 for latest lora support + gpu="H200", + memory=(1024, int(2 * 1024 * 1024)), + image_run_commands=[ + f"rm -rf {HF_CACHE_PATH} 2>/dev/null || true", + + "rm -rf /usr/local/lib/python3.12/dist-packages/nvidia/cudnn/ 2>/dev/null || true", + + "pip install --no-cache-dir hf_xet", + ], + image_env={ + "LD_LIBRARY_PATH": "/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH", + + "HF_XET_HIGH_PERFORMANCE": "1", # for downloading + }, +) + + +class _Miles(MilesConfig): + # Architecture only (MODEL_ARGS); --spec inside it is inert under bridge LoRA. + miles_model_script = "scripts/models/glm5.2-744B-A40B_lora.sh" + + environment = { + "PYTHONPATH": "/root/Megatron-LM/", + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "NCCL_NVLS_ENABLE": "1", + # extra env vars from run_glm5_2_744b_a40b_lora.py + "MILES_EXPERIMENTAL_ROLLOUT_REFACTOR": "1", + "INDEXER_ROPE_NEOX_STYLE": "0", + "SGLANG_NSA_FORCE_MLA": "1", + } + + + hf_checkpoint = "zai-org/GLM-5.2" + megatron_to_hf_mode = "bridge" + + # tilelang + thd backward pass produces nan gradients under nonzero loss (TODO: figure out why?) + # keep megatron for now -- upstream config had data + dsa_attention_backend = "megatron" + qkv_format = "bshd" + data_pad_size_multiplier = 32 + micro_batch_size = 1 + save = f"{CHECKPOINTS_PATH}/GLM-5.2-lora-ckpt" + save_interval = 20 + + + actor_num_nodes = 8 + actor_num_gpus_per_node = 8 + num_gpus_per_node = 8 + colocate = True + use_miles_router = True + calculate_per_token_loss = True + tensor_model_parallel_size = 8 + sequence_parallel = True + pipeline_model_parallel_size = 1 + context_parallel_size = 1 + expert_model_parallel_size = 32 + expert_tensor_parallel_size = 1 + moe_token_dispatcher_type = "alltoall" + + # attention/MLA on every layer + per expert linear_fc1 only on last 10 moe layers + # exclude expert down proj + lora_rank = 8 + lora_alpha = 16 + lora_dropout = 0.0 + target_modules = ( + "q_a_proj,kv_a_proj_with_mqa,q_b_proj,kv_b_proj,o_proj," + "*.layers.68.*.linear_fc1,*.layers.69.*.linear_fc1," + "*.layers.70.*.linear_fc1,*.layers.71.*.linear_fc1," + "*.layers.72.*.linear_fc1,*.layers.73.*.linear_fc1," + "*.layers.74.*.linear_fc1,*.layers.75.*.linear_fc1," + "*.layers.76.*.linear_fc1,*.layers.77.*.linear_fc1" + ) + experts_shared_outer_loras = False + lora_base_cpu_backup = True + no_gradient_accumulation_fusion = True + + + prompt_data = f"{DATA_PATH}/gsm8k/train.parquet" + input_key = "messages" + label_key = "label" + apply_chat_template = True + rollout_shuffle = True + rm_type = "math" + + + num_rollout = 50 + rollout_batch_size = 8 + n_samples_per_prompt = 16 + rollout_max_response_len = 256 + rollout_temperature = 1.0 + global_batch_size = 64 + use_rollout_routing_replay = True + + + advantage_estimator = "grpo" + kl_loss_coef = 0.0 + kl_loss_type = "low_var_kl" + kl_coef = 0.0 + entropy_coef = 0.0 + eps_clip = 0.2 + eps_clip_high = 0.28 + + + optimizer = "adam" + lr = 1e-5 + lr_decay_style = "constant" + weight_decay = 0.1 + adam_beta1 = 0.9 + adam_beta2 = 0.98 + optimizer_cpu_offload = True + overlap_cpu_optimizer_d2h_h2d = True + use_precision_aware_optimizer = True + + + attention_dropout = 0.0 + hidden_dropout = 0.0 + accumulate_allreduce_grads_in_fp32 = True + attention_softmax_in_fp32 = True + attention_backend = "flash" + + # do bf16 sglang rollout -- todo try fp8 rollout and compare logprob diff + rollout_num_gpus_per_engine = 32 + sglang_mem_fraction_static = 0.7 + sglang_enable_dp_attention = True + sglang_ep_size = 32 + sglang_dp_size = 32 + sglang_moe_dense_tp_size = 1 + sglang_enable_dp_lm_head = True + sglang_attention_backend = "nsa" + sglang_nsa_decode_backend = "flashmla_sparse" + sglang_nsa_prefill_backend = "flashmla_sparse" + sglang_page_size = 64 + sglang_cuda_graph_max_bs = 64 + sglang_max_running_requests = 512 + sglang_chunked_prefill_size = 65536 + sglang_watchdog_timeout = 3600 + sglang_moe_runner_backend = "triton" + sglang_disable_shared_experts_fusion = True + sglang_max_lora_rank = 16 + sglang_lora_backend = "triton" + sglang_lora_use_virtual_experts = True + + + use_wandb = True + wandb_project = "miles-run_glm5_2_744b_a40b_lora" + wandb_group = "glm5.2-744B-8node-no-down-proj-megatron-pad32-modal" + disable_wandb_random_suffix = True + + def download_model(self) -> None: + + from huggingface_hub import snapshot_download + + snapshot_download(self.hf_checkpoint, max_workers=32) + + def download_data(self) -> None: + import os + + from huggingface_hub import snapshot_download + + os.makedirs(f"{DATA_PATH}/gsm8k", exist_ok=True) + snapshot_download( + repo_id="zhuzilin/gsm8k", + repo_type="dataset", + local_dir=f"{DATA_PATH}/gsm8k", + ) + + +miles = _Miles() diff --git a/miles/configs/glm5_2_744b_a40b_lora_5layer.py b/miles/configs/glm5_2_744b_a40b_lora_5layer.py new file mode 100644 index 0000000..ebfb0a3 --- /dev/null +++ b/miles/configs/glm5_2_744b_a40b_lora_5layer.py @@ -0,0 +1,152 @@ +"""GLM-5.2 (744B-A40B arch, 5-layer prune) LoRA GRPO — single node, colocated. + +Smoke test for the GLM-5.2 bridge-mode DSA LoRA path. ``Pinaster/GLM-5.2_5layer`` +is a 5-layer prune (3 dense + 2 MoE) of GLM-5.2 that keeps one computing + one +skip layer, so it exercises the same DSA cross-layer index-sharing, MoE, bridge +LoRA, and sglang MoE-LoRA serving path as the full 744B model at toy cost. + +Ports ``scripts/run_glm5_2_744b_a40b_lora.py`` (which the guide launcher does NOT +run directly) into config attributes: the model ``.sh`` supplies architecture +only, every LoRA/DSA/sglang flag is set here and forwarded by ``cli_args()``. + +Requires a miles image built after PR #1559 (GLM-5/5.1/5.2 LoRA) and PR #1593 +(bridge-LoRA recompute fix); the repo default dev-202605291323 predates both. + +Launched by the dedicated smoke harness (pinned to this config): + uv run modal run miles/modal_train_glm_test.py::download_model + uv run modal run miles/modal_train_glm_test.py::download_data + uv run modal run miles/modal_train_glm_test.py::train +""" + +from configs.base import ModalConfig, MilesConfig, DATA_PATH, CHECKPOINTS_PATH, HF_CACHE_PATH + +modal = ModalConfig( + docker_image="radixark/miles:dev-202607090055", + gpu="H200", + memory=(1024, int(2 * 1024 * 1024)), + image_run_commands=[ + f"rm -rf {HF_CACHE_PATH} 2>/dev/null || true", + "rm -rf /usr/local/lib/python3.12/dist-packages/nvidia/cudnn/ 2>/dev/null || true", + ], + image_env={"LD_LIBRARY_PATH": "/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH"}, +) + + +class _Miles(MilesConfig): + miles_model_script = "scripts/models/glm5.2-744B-A40B_5layer_lora.sh" + + environment = { + "PYTHONPATH": "/root/Megatron-LM/", + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "NCCL_NVLS_ENABLE": "1", + "MILES_EXPERIMENTAL_ROLLOUT_REFACTOR": "1", + "INDEXER_ROPE_NEOX_STYLE": "0", + "SGLANG_NSA_FORCE_MLA": "1", + } + + hf_checkpoint = "Pinaster/GLM-5.2_5layer" + megatron_to_hf_mode = "bridge" + dsa_attention_backend = "tilelang" + qkv_format = "thd" + micro_batch_size = 1 + save = f"{CHECKPOINTS_PATH}/GLM-5.2_5layer-lora-ckpt" + save_interval = 1 + + actor_num_nodes = 1 + actor_num_gpus_per_node = 4 + colocate = True + use_miles_router = True + calculate_per_token_loss = True + tensor_model_parallel_size = 4 + sequence_parallel = True + pipeline_model_parallel_size = 1 + context_parallel_size = 1 + expert_model_parallel_size = 4 + expert_tensor_parallel_size = 1 + + lora_rank = 16 + lora_alpha = 32 + lora_dropout = 0.0 + target_modules = "q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj,q_a_proj,kv_a_proj_with_mqa,q_b_proj,kv_b_proj" + experts_shared_outer_loras = True + lora_base_cpu_backup = True + no_gradient_accumulation_fusion = True + + prompt_data = f"{DATA_PATH}/gsm8k/train.parquet" + input_key = "messages" + label_key = "label" + apply_chat_template = True + rollout_shuffle = True + rm_type = "math" + + num_rollout = 1 + rollout_batch_size = 4 + n_samples_per_prompt = 4 + rollout_max_response_len = 512 + rollout_temperature = 1.0 + global_batch_size = 16 + use_rollout_routing_replay = True + + advantage_estimator = "grpo" + kl_loss_coef = 0.0 + kl_loss_type = "low_var_kl" + kl_coef = 0.0 + entropy_coef = 0.0 + eps_clip = 0.2 + eps_clip_high = 0.28 + + optimizer = "adam" + lr = 1e-5 + lr_decay_style = "constant" + weight_decay = 0.1 + adam_beta1 = 0.9 + adam_beta2 = 0.98 + optimizer_cpu_offload = True + overlap_cpu_optimizer_d2h_h2d = True + use_precision_aware_optimizer = True + + attention_dropout = 0.0 + hidden_dropout = 0.0 + accumulate_allreduce_grads_in_fp32 = True + attention_softmax_in_fp32 = True + attention_backend = "flash" + + rollout_num_gpus_per_engine = 2 + sglang_mem_fraction_static = 0.5 + sglang_enable_dp_attention = True + sglang_ep_size = 2 + sglang_dp_size = 2 + sglang_moe_dense_tp_size = 1 + sglang_enable_dp_lm_head = True + sglang_attention_backend = "nsa" + sglang_nsa_decode_backend = "flashmla_sparse" + sglang_nsa_prefill_backend = "flashmla_sparse" + sglang_page_size = 64 + sglang_cuda_graph_max_bs = 64 + sglang_max_running_requests = 512 + sglang_chunked_prefill_size = 4096 + sglang_watchdog_timeout = 3600 + sglang_moe_runner_backend = "triton" + sglang_disable_shared_experts_fusion = True + sglang_max_lora_rank = 16 + sglang_lora_backend = "triton" + + use_wandb = True + wandb_project = "miles-run_glm5_2_744b_a40b_lora" + wandb_group = "glm5.2-5layer-lora" + disable_wandb_random_suffix = True + + def download_data(self) -> None: + import os + + from huggingface_hub import snapshot_download + + os.makedirs(f"{DATA_PATH}/gsm8k", exist_ok=True) + snapshot_download( + repo_id="zhuzilin/gsm8k", + repo_type="dataset", + local_dir=f"{DATA_PATH}/gsm8k", + ) + + +miles = _Miles() diff --git a/miles/configs/glm5_2_744b_a40b_lora_dapo.py b/miles/configs/glm5_2_744b_a40b_lora_dapo.py new file mode 100644 index 0000000..40bd7bf --- /dev/null +++ b/miles/configs/glm5_2_744b_a40b_lora_dapo.py @@ -0,0 +1,155 @@ +"""GLM-5.2 (full 744B-A40B) LoRA GRPO on dapo-math — long-context DSA variant. + +untested -- same recipe as glm5_2_744b_a40b_lora.py but on dapo-math task instead of gsm8k and increase +rollout respone length to 4096 tokens and context window to 8192 tokens + + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora_dapo uv run modal run miles/modal_train.py::download_model + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora_dapo uv run modal run miles/modal_train.py::download_data + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora_dapo uv run modal run miles/modal_train.py::train + +""" + +from configs.base import ModalConfig, MilesConfig, DATA_PATH, CHECKPOINTS_PATH, HF_CACHE_PATH + +modal = ModalConfig( + docker_image="radixark/miles:dev-202607090055", + gpu="H200", + memory=(1024, int(2 * 1024 * 1024)), + image_run_commands=[ + f"rm -rf {HF_CACHE_PATH} 2>/dev/null || true", + "rm -rf /usr/local/lib/python3.12/dist-packages/nvidia/cudnn/ 2>/dev/null || true", + "pip install --no-cache-dir hf_xet", + ], + image_env={ + "LD_LIBRARY_PATH": "/usr/lib/x86_64-linux-gnu:$LD_LIBRARY_PATH", + "HF_XET_HIGH_PERFORMANCE": "1", + }, +) + + +class _Miles(MilesConfig): + miles_model_script = "scripts/models/glm5.2-744B-A40B_lora.sh" + + environment = { + "PYTHONPATH": "/root/Megatron-LM/", + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "NCCL_NVLS_ENABLE": "1", + "MILES_EXPERIMENTAL_ROLLOUT_REFACTOR": "1", + "INDEXER_ROPE_NEOX_STYLE": "0", + "SGLANG_NSA_FORCE_MLA": "1", + } + + hf_checkpoint = "zai-org/GLM-5.2" + megatron_to_hf_mode = "bridge" + dsa_attention_backend = "tilelang" + qkv_format = "thd" + micro_batch_size = 1 + save = f"{CHECKPOINTS_PATH}/GLM-5.2-lora-dapo-ckpt" + save_interval = 1 + + actor_num_nodes = 1 + actor_num_gpus_per_node = 8 + colocate = True + use_miles_router = True + calculate_per_token_loss = True + tensor_model_parallel_size = 8 + sequence_parallel = True + pipeline_model_parallel_size = 1 + context_parallel_size = 1 + expert_model_parallel_size = 8 + expert_tensor_parallel_size = 1 + + lora_rank = 16 + lora_alpha = 32 + lora_dropout = 0.0 + target_modules = "q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj,q_a_proj,kv_a_proj_with_mqa,q_b_proj,kv_b_proj" + experts_shared_outer_loras = True + lora_base_cpu_backup = True + no_gradient_accumulation_fusion = True + + prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" + input_key = "prompt" + label_key = "label" + apply_chat_template = True + rollout_shuffle = True + rm_type = "math" + + num_rollout = 1 + rollout_batch_size = 4 + n_samples_per_prompt = 4 + rollout_max_response_len = 4096 + seq_length = 8192 + rollout_max_context_len = 8192 + rollout_temperature = 1.0 + global_batch_size = 16 + use_rollout_routing_replay = True + + advantage_estimator = "grpo" + kl_loss_coef = 0.0 + kl_loss_type = "low_var_kl" + kl_coef = 0.0 + entropy_coef = 0.0 + eps_clip = 0.2 + eps_clip_high = 0.28 + + optimizer = "adam" + lr = 1e-5 + lr_decay_style = "constant" + weight_decay = 0.1 + adam_beta1 = 0.9 + adam_beta2 = 0.98 + optimizer_cpu_offload = True + overlap_cpu_optimizer_d2h_h2d = True + use_precision_aware_optimizer = True + + attention_dropout = 0.0 + hidden_dropout = 0.0 + accumulate_allreduce_grads_in_fp32 = True + attention_softmax_in_fp32 = True + attention_backend = "flash" + + rollout_num_gpus_per_engine = 8 + sglang_mem_fraction_static = 0.5 + sglang_enable_dp_attention = True + sglang_ep_size = 8 + sglang_dp_size = 8 + sglang_moe_dense_tp_size = 1 + sglang_enable_dp_lm_head = True + sglang_attention_backend = "nsa" + sglang_nsa_decode_backend = "flashmla_kv" + sglang_nsa_prefill_backend = "flashmla_sparse" + sglang_kv_cache_dtype = "fp8_e4m3" + sglang_page_size = 64 + sglang_cuda_graph_max_bs = 256 + sglang_max_running_requests = 512 + sglang_chunked_prefill_size = 16384 + sglang_watchdog_timeout = 3600 + sglang_moe_runner_backend = "triton" + sglang_disable_shared_experts_fusion = True + sglang_max_lora_rank = 16 + sglang_lora_backend = "triton" + + use_wandb = True + wandb_project = "miles-run_glm5_2_744b_a40b_lora" + wandb_group = "glm5.2-744B-lora-dapo" + disable_wandb_random_suffix = True + + def download_model(self) -> None: + from huggingface_hub import snapshot_download + + snapshot_download(self.hf_checkpoint, max_workers=32) + + def download_data(self) -> None: + import os + + from huggingface_hub import snapshot_download + + os.makedirs(f"{DATA_PATH}/dapo-math-17k", exist_ok=True) + snapshot_download( + repo_id="zhuzilin/dapo-math-17k", + repo_type="dataset", + local_dir=f"{DATA_PATH}/dapo-math-17k", + ) + + +miles = _Miles() diff --git a/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k.py b/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k.py new file mode 100644 index 0000000..b039559 --- /dev/null +++ b/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k.py @@ -0,0 +1,92 @@ +"""GLM-5.2 full-model TileLang/THD on dapo-math — 32k context + 1k rollout, 8 nodes x 8 H200. + +Uses the validated 64-GPU TileLang setup (safe-indices kernel patch, full +activation recompute, extended SGLang load timeout) with the dapo-math task. +No context parallelism (CP=1); 32k context fits without it (validated 2026-07-12). + + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora_dapo_tilelang_32k \ + uv run modal run miles/modal_train.py::download_data + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora_dapo_tilelang_32k \ + uv run modal run --detach miles/modal_train.py::train +""" + +from configs import glm5_2_744b_a40b_lora as _base +from configs.base import CHECKPOINTS_PATH, DATA_PATH, ModalConfig + +_PATCH = "miles/glm5_tilelang_safe_indices.patch" + +modal = ModalConfig( + docker_image=_base.modal.docker_image, + gpu=_base.modal.gpu, + memory=_base.modal.memory, + cloud=_base.modal.cloud, + region=_base.modal.region, + patch_files=[_PATCH], + image_run_commands=[ + *_base.modal.image_run_commands, + ( + "cd /usr/local/lib/python3.12/dist-packages && " + "git apply --check /tmp/glm5_tilelang_safe_indices.patch && " + "git apply /tmp/glm5_tilelang_safe_indices.patch" + ), + ( + "python -c \"from pathlib import Path; " + "p = Path('/sgl-workspace/sglang/python/sglang/srt/model_executor/model_runner.py'); " + "s = p.read_text(); " + "old = 'UNBALANCED_MODEL_LOADING_TIMEOUT_S = 480'; " + "assert s.count(old) == 1; " + "p.write_text(s.replace(old, 'UNBALANCED_MODEL_LOADING_TIMEOUT_S = 1800'))\"" + ), + # LoRA checkpoint fix: dp_rank_0's mkdir on the shared volume is not + # visible to containers on other nodes, so every rank must mkdir itself + # before writing its training_state_rank{N}.pt. + ( + "python -c \"from pathlib import Path; " + "p = Path('/root/miles/miles/backends/megatron_utils/lora_utils.py'); " + "s = p.read_text(); " + "old = 'if is_dp_rank_0:\\n save_path.mkdir(parents=True, exist_ok=True)'; " + "new = 'save_path.mkdir(parents=True, exist_ok=True)'; " + "assert s.count(old) == 1; " + "p.write_text(s.replace(old, new))\"" + ), + ], + image_env=dict(_base.modal.image_env), +) + + +class _Miles(_base._Miles): + dsa_attention_backend = "tilelang" + qkv_format = "thd" + data_pad_size_multiplier = None + recompute_granularity = "full" + recompute_method = "uniform" + recompute_num_layers = 1 + + prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" + input_key = "prompt" + label_key = "label" + + seq_length = 32768 + rollout_max_context_len = 32768 + rollout_max_response_len = 1024 + + num_rollout = 50 + save = f"{CHECKPOINTS_PATH}/GLM-5.2-lora-dapo-32k-ckpt" + save_interval = 10 + + wandb_group = "glm5.2-744B-8node-tilelang-dapo-32k" + + def download_data(self) -> None: + import os + + from huggingface_hub import snapshot_download + + os.makedirs(f"{DATA_PATH}/dapo-math-17k", exist_ok=True) + snapshot_download( + repo_id="zhuzilin/dapo-math-17k", + repo_type="dataset", + local_dir=f"{DATA_PATH}/dapo-math-17k", + ) + + +miles = _Miles() diff --git a/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_16node.py b/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_16node.py new file mode 100644 index 0000000..9e5933c --- /dev/null +++ b/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_16node.py @@ -0,0 +1,113 @@ +"""GLM-5.2 LoRA DAPO 32k padded — 16 nodes, NO activation recompute, bf16 rollouts. +tp 8 x cp 4 -> 32 gpu replicas, dp 4, ep 32 unchanged + +""" + +from configs import glm5_2_744b_a40b_lora_dapo_tilelang_32k as _base +from configs.base import CHECKPOINTS_PATH, DATA_PATH, ModalConfig + +modal = ModalConfig( + docker_image=_base.modal.docker_image, + gpu=_base.modal.gpu, + memory=_base.modal.memory, + cloud=_base.modal.cloud, + region=_base.modal.region, + patch_files=[*_base.modal.patch_files, "miles/megatron_dsa_cp_assert_fix.py"], + image_run_commands=[ + *_base.modal.image_run_commands, + # megatron-core blanket-refuses CP for DSA; the TileLang bridge path + # has its own CP collectives (see the patch docstring), so gate the + # assert on the backend instead. + "python /tmp/megatron_dsa_cp_assert_fix.py", + ], + image_env=dict(_base.modal.image_env), +) + + +# filler tokens to mimic 32k context length +# 26k filler + header + question (~0.3-1.5k) ≈ 26.5-27.5k prompt tokens, +# + 4096 response ≤ ~31.7k, under the 32767 limit. +_FILLER_TOKENS = 26000 + + +class _Miles(_base._Miles): + actor_num_nodes = 16 + + # TP8 x CP4 -> 32-GPU replicas, DP 4. EP32 unchanged (128 % 32 == 0). + context_parallel_size = 4 + allgather_cp = True + + # The whole point of this config: no activation recompute. + recompute_granularity = None + recompute_method = None + recompute_num_layers = None + + prompt_data = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k-pad26k.jsonl" + rollout_max_response_len = 4096 + + # 2 gradient steps per rollout (128 samples / global_batch_size 64). + num_rollout = 40 + save = f"{CHECKPOINTS_PATH}/GLM-5.2-lora-dapo-32k-16node-ckpt" + save_interval = 10 + + wandb_group = "glm5.2-744B-16node-tilelang-dapo-32k-norecompute" + + def download_data(self) -> None: + """Generate the padded dapo set (same as the retired padstress config).""" + import json + import os + import random + + from huggingface_hub import snapshot_download + from transformers import AutoTokenizer + + os.makedirs(f"{DATA_PATH}/dapo-math-17k", exist_ok=True) + snapshot_download( + repo_id="zhuzilin/dapo-math-17k", + repo_type="dataset", + local_dir=f"{DATA_PATH}/dapo-math-17k", + ) + + dst = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k-pad26k.jsonl" + if os.path.exists(dst): + print(f"{dst} already exists, skipping generation") + return + + tokenizer = AutoTokenizer.from_pretrained( + "zai-org/GLM-5.2", trust_remote_code=True + ) + words = ( + "system model tensor kernel matrix vector gradient layer token " + "attention memory cache buffer stream block thread warp shard " + "sequence batch epoch metric loss reward policy value state action" + ).split() + + + rng = random.Random(0) + base = " ".join(rng.choice(words) for _ in range(2 * _FILLER_TOKENS)) + filler = tokenizer.decode( + tokenizer.encode(base, add_special_tokens=False)[:_FILLER_TOKENS] + ) + + src = f"{DATA_PATH}/dapo-math-17k/dapo-math-17k.jsonl" + n_written = 0 + with open(src) as fin, open(dst, "w") as fout: + for i, line in enumerate(fin): + sample = json.loads(line) + question = sample["prompt"][0]["content"] + sample["prompt"][0]["content"] = ( + f"Reference log #{i}-{random.Random(i).getrandbits(64):x} " + "follows. It is not relevant to the question; ignore it and " + "solve the question at the end.\n\n" + f"{filler}\n\n{question}" + ) + fout.write(json.dumps(sample) + "\n") + n_written += 1 + + total = len( + tokenizer.encode(json.loads(open(dst).readline())["prompt"][0]["content"]) + ) + print(f"Wrote {n_written} padded samples to {dst}; sample 0 = {total} tokens") + + +miles = _Miles() diff --git a/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_8node_cp4.py b/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_8node_cp4.py new file mode 100644 index 0000000..9d12c24 --- /dev/null +++ b/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_8node_cp4.py @@ -0,0 +1,29 @@ +"""GLM-5.2 LoRA DAPO 32k padded — 8 nodes, CP4, NO activation recompute. + +Same TP8 x CP4 geometry (and hence identical per-rank memory picture) as the +16-node config, on the readily-schedulable 8-node block: replica = 32 GPUs, +DP = 2. See glm5_2_744b_a40b_lora_dapo_tilelang_32k_16node for the CP design +notes and the megatron_dsa_cp_assert_fix rationale. + +Correctness oracle: train_rollout_logprob_abs_diff (~0.01 healthy, >1 means +broken CP attention math — stop the run if seen). + + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora_dapo_tilelang_32k_8node_cp4 \ + uv run modal run --detach miles/modal_train.py::train +""" + +from configs import glm5_2_744b_a40b_lora_dapo_tilelang_32k_16node as _base +from configs.base import CHECKPOINTS_PATH + +modal = _base.modal + + +class _Miles(_base._Miles): + actor_num_nodes = 8 + + save = f"{CHECKPOINTS_PATH}/GLM-5.2-lora-dapo-32k-8node-cp4-ckpt" + + wandb_group = "glm5.2-744B-8node-cp4-tilelang-dapo-32k-norecompute" + + +miles = _Miles() diff --git a/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8.py b/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8.py new file mode 100644 index 0000000..5ec0f86 --- /dev/null +++ b/miles/configs/glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8.py @@ -0,0 +1,101 @@ +"""GLM-5.2 LoRA DAPO training with FP8 rollout — 8 nodes x 8 H200, TileLang/THD, 32k context. + +The finalized FP8-rollout setup: the rollout engines serve the official +block-quantized ``zai-org/GLM-5.2-FP8`` checkpoint while the trainer keeps +bf16 ``zai-org/GLM-5.2``. LoRA adapters (bf16) sync to the engines every +rollout. Measured vs the bf16/bf16 baseline: generation ~25-40% faster +(351-450s vs ~540s per rollout), train_rollout_logprob_abs_diff ~0.040 +(bf16 ~0.010, online-quant fp8 ~0.057), no sampling NaNs, rewards learn. + +Fixes this configuration depends on (see each patch's docstring): + - glm5_tilelang_safe_indices.patch: TileLang sparse-MLA backward NaNs + (unsafe padded-index access + aggressive shared-memory merge miscompile). + - sglang_fp8_lora_fix.py: LoRA-B buffer mis-sizing on quantized + column-parallel layers crashed engine init under --quantization fp8 / + quantized checkpoints. + - sglang_tp1_shared_expert_fix.py: SGLANG_SHARED_EXPERT_TP1-replicated + shared expert was double-added (once per TP rank) whenever the post-MoE + all-reduce is deferred/replaced (FlashInfer AllReduce Fusion, + dp-attention reduce-scatterv); folded into the pre-reduction add. + - SGLANG_SHARED_EXPERT_TP1=1: the 128x128 block-quantized checkpoint cannot + TP32-shard the shared expert (2048/32 = 64-row shards < one scale block), + so replicate it instead (~25 MB/rank). + + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8 \ + uv run modal run miles/modal_train.py::download_model + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8 \ + uv run modal run miles/modal_train.py::download_data + EXPERIMENT_CONFIG=glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8 \ + uv run modal run --detach miles/modal_train.py::train +""" + +from configs import glm5_2_744b_a40b_lora_dapo_tilelang_32k as _base +from configs.base import CHECKPOINTS_PATH, ModalConfig + +_FP8_CHECKPOINT = "zai-org/GLM-5.2-FP8" + +modal = ModalConfig( + docker_image=_base.modal.docker_image, + gpu=_base.modal.gpu, + memory=_base.modal.memory, + cloud=_base.modal.cloud, + region=_base.modal.region, + patch_files=[ + *_base.modal.patch_files, + "miles/sglang_fp8_lora_fix.py", + "miles/sglang_tp1_shared_expert_fix.py", + ], + image_run_commands=[ + *_base.modal.image_run_commands, + "python /tmp/sglang_fp8_lora_fix.py", + "python /tmp/sglang_tp1_shared_expert_fix.py", + ], + image_env=dict(_base.modal.image_env), +) + + +class _Miles(_base._Miles): + environment = { + **_base._Miles.environment, + "SGLANG_SHARED_EXPERT_TP1": "1", + } + + # The pre-quantized checkpoint carries its own quantization_config; + # do NOT also force online quantization. + sglang_quantization = None + + # Rollout base differs from hf_checkpoint, so adapter sync must be + # opted into explicitly. + sglang_config = { + "sglang": [ + { + "name": "actor", + "model_path": _FP8_CHECKPOINT, + "update_weights": True, + "num_gpus_per_engine": 32, + "server_groups": [ + {"worker_type": "regular", "num_gpus": 64}, + ], + } + ] + } + + # 4096-token response budget so completions can finish and reward can + # improve (1024 truncated heavily on dapo-math). + rollout_max_response_len = 4096 + + # 2 gradient steps per rollout (128 samples / global_batch_size 64). + num_rollout = 50 + save = f"{CHECKPOINTS_PATH}/GLM-5.2-lora-dapo-32k-fp8-ckpt" + save_interval = 10 + + wandb_group = "glm5.2-744B-8node-tilelang-dapo-32k-fp8" + + def download_model(self) -> None: + from huggingface_hub import snapshot_download + + snapshot_download(self.hf_checkpoint, max_workers=32) + snapshot_download(_FP8_CHECKPOINT, max_workers=32) + + +miles = _Miles() diff --git a/miles/glm5_tilelang_safe_indices.patch b/miles/glm5_tilelang_safe_indices.patch new file mode 100644 index 0000000..4337905 --- /dev/null +++ b/miles/glm5_tilelang_safe_indices.patch @@ -0,0 +1,157 @@ +diff --git a/megatron/bridge/models/glm5/tilelang/sparse_mla.py b/megatron/bridge/models/glm5/tilelang/sparse_mla.py +--- a/megatron/bridge/models/glm5/tilelang/sparse_mla.py ++++ b/megatron/bridge/models/glm5/tilelang/sparse_mla.py +@@ -59,6 +59,14 @@ class SparseMLA(torch.autograd.Function): + scaling = ctx.scaling + + tl_dq, tl_dkv = sparse_mla_bwd(q, kv, tl_out, grad_output.contiguous(), indices, tl_lse, sm_scale=scaling) ++ for name, grad in (("dQ", tl_dq), ("dKV", tl_dkv)): ++ finite = torch.isfinite(grad) ++ if not finite.all(): ++ nonfinite = (~finite).sum().item() ++ raise FloatingPointError( ++ f"TileLang SparseMLA {name} contains {nonfinite} non-finite values " ++ f"(shape={tuple(grad.shape)}, dtype={grad.dtype})" ++ ) + + # Return gradients for each input (None for indices as it's not differentiable) + return tl_dq, tl_dkv, None, None +diff --git a/megatron/bridge/models/glm5/tilelang/tilelang_sparse_mla_fwd.py b/megatron/bridge/models/glm5/tilelang/tilelang_sparse_mla_fwd.py +--- a/megatron/bridge/models/glm5/tilelang/tilelang_sparse_mla_fwd.py ++++ b/megatron/bridge/models/glm5/tilelang/tilelang_sparse_mla_fwd.py +@@ -109,6 +109,7 @@ def sparse_mla_fwd( + O_shared = T.alloc_shared([H_per_block, D], dtype) + Lse_shared = T.alloc_shared([H_per_block], accum_dtype) + mask = T.alloc_fragment([BI], "bool") ++ safe_indices = T.alloc_shared([BI], indices_dtype) + + acc_o = T.alloc_fragment([H_per_block, D], accum_dtype) + acc_s = T.alloc_fragment([H_per_block, BI], accum_dtype) +@@ -136,13 +137,25 @@ def sparse_mla_fwd( + + for i_i in T.Pipelined(NI, num_stages=num_stages): + for bi_i in T.Parallel(BI): +- # Changed here for thd +- mask[bi_i] = Indices[b_i, s_i, g_i, i_i * BI + bi_i] != -1 ++ mask[bi_i] = ( ++ Indices[b_i, s_i, g_i, i_i * BI + bi_i] >= 0 ++ and Indices[b_i, s_i, g_i, i_i * BI + bi_i] < seq_len_kv ++ ) ++ safe_indices[bi_i] = T.if_then_else( ++ mask[bi_i], ++ Indices[b_i, s_i, g_i, i_i * BI + bi_i], ++ 0, ++ ) ++ T.sync_threads() + + for bi_i, d_i in T.Parallel(BI, D): +- KV_shared[bi_i, d_i] = KV[b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, d_i] ++ KV_shared[bi_i, d_i] = T.if_then_else( ++ mask[bi_i], KV[b_i, safe_indices[bi_i], g_i, d_i], 0 ++ ) + for bi_i, d_i in T.Parallel(BI, D_tail): +- K_tail_shared[bi_i, d_i] = KV[b_i, Indices[b_i, s_i, g_i, i_i * BI + bi_i], g_i, D + d_i] ++ K_tail_shared[bi_i, d_i] = T.if_then_else( ++ mask[bi_i], KV[b_i, safe_indices[bi_i], g_i, D + d_i], 0 ++ ) + + for h_i, bi_i in T.Parallel(H_per_block, BI): + acc_s[h_i, bi_i] = T.if_then_else(mask[bi_i], 0, -T.infinity(acc_s.dtype)) +diff --git a/megatron/bridge/models/glm5/tilelang/tilelang_sparse_mla_bwd.py b/megatron/bridge/models/glm5/tilelang/tilelang_sparse_mla_bwd.py +--- a/megatron/bridge/models/glm5/tilelang/tilelang_sparse_mla_bwd.py ++++ b/megatron/bridge/models/glm5/tilelang/tilelang_sparse_mla_bwd.py +@@ -98,7 +98,7 @@ def bwd( + pass_configs={ + tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True, + tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True, +- tilelang.PassConfigKey.TL_ENABLE_AGGRESSIVE_SHARED_MEMORY_MERGE: True, ++ tilelang.PassConfigKey.TL_ENABLE_AGGRESSIVE_SHARED_MEMORY_MERGE: False, + }, + ) + def bwd( +@@ -168,6 +168,7 @@ def bwd( + KV_tail_shared = T.alloc_shared([BS, D_tail], dtype) + dO_shared = T.alloc_shared([block_H, D], dtype) + mask = T.alloc_fragment([BS], "bool") ++ safe_indices = T.alloc_shared([BS], indices_dtype) + + P_shared_cast = T.alloc_shared([block_H, BS], dtype) + dP_shared_cast = T.alloc_shared([block_H, BS], dtype) +@@ -195,21 +196,33 @@ def bwd( + # Process each block of indices + for i_i in T.Pipelined(NS, num_stages=num_stages): + # Check which indices are valid + for bi_i in T.Parallel(BS): +- # Changed here for thd +- mask[bi_i] = Indices[by, s_i, bz // NH, i_i * BS + bi_i] != -1 ++ mask[bi_i] = ( ++ Indices[by, s_i, bz // NH, i_i * BS + bi_i] >= 0 ++ and Indices[by, s_i, bz // NH, i_i * BS + bi_i] < S_kv ++ ) ++ safe_indices[bi_i] = T.if_then_else( ++ mask[bi_i], ++ Indices[by, s_i, bz // NH, i_i * BS + bi_i], ++ 0, ++ ) ++ T.sync_threads() + + # Compute attention scores + for h_i, bi_i in T.Parallel(block_H, BS): + acc_p[h_i, bi_i] = T.if_then_else(mask[bi_i], 0, -T.infinity(acc_p.dtype)) + + # Load KV, V for this block of indices + for bi_i, d_i in T.Parallel(BS, D): +- KV_shared[bi_i, d_i] = KV[by, Indices[by, s_i, bz // NH, i_i * BS + bi_i], bz // NH, d_i] ++ KV_shared[bi_i, d_i] = T.if_then_else( ++ mask[bi_i], KV[by, safe_indices[bi_i], bz // NH, d_i], 0 ++ ) + + T.gemm(Q_shared, KV_shared, acc_p, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) + + for bi_i, d_i in T.Parallel(BS, D_tail): +- KV_tail_shared[bi_i, d_i] = KV[by, Indices[by, s_i, bz // NH, i_i * BS + bi_i], bz // NH, D + d_i] ++ KV_tail_shared[bi_i, d_i] = T.if_then_else( ++ mask[bi_i], KV[by, safe_indices[bi_i], bz // NH, D + d_i], 0 ++ ) + T.gemm(Q_tail_shared, KV_tail_shared, acc_p, transpose_B=True, policy=T.GemmWarpPolicy.FullCol) + +@@ -256,27 +269,21 @@ def bwd( + acc_dkv_tail_shared[bi_i, d_i] = acc_dkv_tail[bi_i + s * (BS // split_store), d_i] + + for bi_i, d_i in T.Parallel(BS // split_store, D // 4): +- T.atomic_addx4( +- dKV[ +- by, +- Indices[by, s_i, bz // NH, i_i * BS + bi_i + s * (BS // split_store)], +- bz // NH, +- d_i * 4, +- ], +- acc_dkv_shared[bi_i, d_i * 4], +- ) ++ index_offset = bi_i + s * (BS // split_store) ++ if mask[index_offset]: ++ T.atomic_addx4( ++ dKV[by, safe_indices[index_offset], bz // NH, d_i * 4], ++ acc_dkv_shared[bi_i, d_i * 4], ++ ) + + # Atomically update dKV, dKV_tail tensors + for bi_i, d_i in T.Parallel(BS // split_store, D_tail // 4): +- T.atomic_addx4( +- dKV[ +- by, +- Indices[by, s_i, bz // NH, i_i * BS + bi_i + s * (BS // split_store)], +- bz // NH, +- D + d_i * 4, +- ], +- acc_dkv_tail_shared[bi_i, d_i * 4], +- ) ++ index_offset = bi_i + s * (BS // split_store) ++ if mask[index_offset]: ++ T.atomic_addx4( ++ dKV[by, safe_indices[index_offset], bz // NH, D + d_i * 4], ++ acc_dkv_tail_shared[bi_i, d_i * 4], ++ ) + + # Store the accumulated dQ + T.copy(acc_dq, dQ_shared) diff --git a/miles/megatron_dsa_cp_assert_fix.py b/miles/megatron_dsa_cp_assert_fix.py new file mode 100644 index 0000000..2798808 --- /dev/null +++ b/miles/megatron_dsa_cp_assert_fix.py @@ -0,0 +1,45 @@ +"""Relax megatron-core's blanket "no CP for DSA" assert for the TileLang backend. + +megatron-core's MCoreMLATransformerConfig.__post_init__ refuses +context_parallel_size > 1 whenever experimental_attention_variant == "dsa". +That is correct for the generic megatron-core DSAttention module (it has no +CP collectives), but Megatron-Bridge's GLM-5 TileLang path replaces that +module with CP-capable code (see bridge models/glm5/tilelang/tilelang_mla.py: +CP-gathered K, CP-local q RoPE replicate/slice, indexer varlen bounds +scattered over the CP group — the slime/baseten allgather-CP scheme). + +This patch gates the assert on the DSA kernel backend: CP stays forbidden +for the megatron backend, and is allowed for tilelang. finalize() calls +__post_init__ on the provider instance, which carries dsa_attention_backend, +so the getattr sees it; the default keeps the assert enforced otherwise. + +Correctness oracle for runs using this: train_rollout_logprob_abs_diff. +The trainer rescoring of SGLang-generated tokens catches wrong CP attention +math immediately (~0.01 healthy vs >1 broken). +""" + +from pathlib import Path + +TARGET = Path("/root/Megatron-LM/megatron/core/transformer/transformer_config.py") + +OLD = ''' elif self.experimental_attention_variant == "dsa": + assert ( + self.context_parallel_size == 1 + ), "Currently context parallelism is not supported by DSAttention!"''' + +NEW = ''' elif self.experimental_attention_variant == "dsa": + assert ( + self.context_parallel_size == 1 + or getattr(self, "dsa_attention_backend", "megatron") == "tilelang" + ), ( + "Context parallelism with DSA requires the TileLang backend " + "(megatron-core DSAttention has no CP collectives)." + )''' + +src = TARGET.read_text() +if NEW in src: + print("already patched, skipping") +else: + assert src.count(OLD) == 1, f"expected 1 match, got {src.count(OLD)}" + TARGET.write_text(src.replace(OLD, NEW)) + print(f"patched {TARGET}: DSA CP assert gated on tilelang backend") diff --git a/miles/modal_sglang_serve_test.py b/miles/modal_sglang_serve_test.py new file mode 100644 index 0000000..4a8a723 --- /dev/null +++ b/miles/modal_sglang_serve_test.py @@ -0,0 +1,164 @@ +"""Standalone 1-node SGLang serving test for GLM-5.2 FP8 debugging. + +Serves the model on 8xH200 with the same SGLang flags as the miles rollout +engines (minus dp-attention, which needs 32 GPUs), then runs greedy +completions with logprobs to judge output quality directly. Bisect levers +are exposed via --server-args / --env-json so variants (shared-expert TP1, +moe runner, nsa backends, online vs offline quant) don't need code changes. + + uv run modal run --detach miles/modal_sglang_serve_test.py::serve_test + uv run modal run --detach miles/modal_sglang_serve_test.py::serve_test \ + --env-json '{"SGLANG_SHARED_EXPERT_TP1": "1"}' +""" + +import json +import os +import subprocess +import time +import urllib.request + +import modal + +from configs import get_module +from configs.base import HF_CACHE_PATH + +# Reuse the FP8 experiment's image (same sglang build + patches as the +# 64-GPU rollout engines, incl. sglang_fp8_lora_fix and the load-timeout bump). +EXPERIMENT = "glm5_2_744b_a40b_lora_dapo_tilelang_32k_fp8" +modal_cfg = get_module(EXPERIMENT).modal + +image = ( + modal.Image.from_registry(modal_cfg.docker_image) + .entrypoint([]) + .add_local_python_source("configs", copy=True) + .add_local_python_source("modal_helpers", copy=True) +) +for patch in modal_cfg.patch_files: + image = image.add_local_file(patch, f"/tmp/{os.path.basename(patch)}", copy=True) +if modal_cfg.image_run_commands: + image = image.run_commands(*modal_cfg.image_run_commands) +if modal_cfg.image_env: + image = image.env(modal_cfg.image_env) + +hf_cache_volume = modal.Volume.from_name("huggingface-cache", create_if_missing=True) + +app = modal.App("glm52-fp8-serve-test") + +PORT = 30000 + +DEFAULT_SERVER_ARGS = ( + "--model-path zai-org/GLM-5.2-FP8 " + "--tp-size 8 " + "--trust-remote-code " + "--attention-backend nsa " + "--nsa-decode-backend flashmla_sparse " + "--nsa-prefill-backend flashmla_sparse " + "--moe-runner-backend triton " + "--disable-shared-experts-fusion " + "--mem-fraction-static 0.80 " + "--context-length 8192 " + # Quality test only: skip graph capture (15 min) and its NVLS multicast + # setup, which Fabric Manager on these hosts cannot provide. + "--disable-cuda-graph " + f"--port {PORT} --host 127.0.0.1" +) + +# Same rollout-engine env as the miles config. +DEFAULT_ENV = { + "SGLANG_NSA_FORCE_MLA": "1", + "INDEXER_ROPE_NEOX_STYLE": "0", + "NCCL_NVLS_ENABLE": "0", +} + +PROMPTS = [ + "The capital of France is", + ( + "Question: Natalia sold clips to 48 of her friends in April, and then " + "she sold half as many clips in May. How many clips did Natalia sell " + "altogether in April and May?\nAnswer:" + ), + "def fibonacci(n):\n", + "1 + 1 = 2, 2 + 2 = 4, 4 + 4 =", +] + + +def _generate(prompt: str, max_new_tokens: int = 96) -> dict: + payload = { + "text": prompt, + "sampling_params": {"temperature": 0, "max_new_tokens": max_new_tokens}, + "return_logprob": True, + } + req = urllib.request.Request( + f"http://127.0.0.1:{PORT}/generate", + data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json"}, + ) + return json.loads(urllib.request.urlopen(req, timeout=900).read()) + + +def _post(path: str, payload: dict | None = None) -> str: + req = urllib.request.Request( + f"http://127.0.0.1:{PORT}/{path}", + data=json.dumps(payload or {}).encode(), + headers={"Content-Type": "application/json"}, + ) + return urllib.request.urlopen(req, timeout=900).read().decode() + + +def _run_prompts(phase: str) -> None: + for prompt in PROMPTS: + out = _generate(prompt) + lps = [t[0] for t in out["meta_info"]["output_token_logprobs"]] + mean_lp = sum(lps) / max(len(lps), 1) + print("=" * 60) + print(f"[{phase}] PROMPT: {prompt[:120]!r}") + print(f"[{phase}] OUTPUT: {out['text'][:400]!r}") + print(f"[{phase}] greedy mean output logprob: {mean_lp:.3f} over {len(lps)} tokens") + + +@app.function( + image=image, + gpu="H200:8", + volumes={str(HF_CACHE_PATH): hf_cache_volume}, + timeout=3 * 60 * 60, + secrets=[modal.Secret.from_name("huggingface-secret")], +) +def serve_test(server_args: str = "", env_json: str = "{}", cycle_memory: bool = False): + hf_cache_volume.reload() + args = server_args or DEFAULT_SERVER_ARGS + if cycle_memory: + # Mirror the miles rollout engines' memory-saver setup so we can + # exercise the release/resume weight backup path the training loop + # performs before rollout 0. + args += " --enable-memory-saver --enable-weights-cpu-backup" + env = {**os.environ, **DEFAULT_ENV, **json.loads(env_json)} + print(f"server args: {args}") + print(f"env overrides: {json.loads(env_json)}") + + proc = subprocess.Popen(f"python -m sglang.launch_server {args}", shell=True, env=env) + try: + deadline = time.time() + 45 * 60 + while True: + if proc.poll() is not None: + raise RuntimeError(f"server exited during startup: {proc.returncode}") + try: + urllib.request.urlopen(f"http://127.0.0.1:{PORT}/health_generate", timeout=5) + break + except Exception: + if time.time() > deadline: + raise TimeoutError("server did not become healthy in 45 min") + time.sleep(10) + print("server healthy, running prompts") + + _run_prompts("fresh") + + if cycle_memory: + print("cycling memory occupation (release -> resume)...") + print(_post("release_memory_occupation")[:200]) + time.sleep(10) + print(_post("resume_memory_occupation")[:200]) + _run_prompts("after-cycle") + + print("DONE_SERVE_TEST") + finally: + proc.terminate() diff --git a/miles/modal_train_glm_test.py b/miles/modal_train_glm_test.py new file mode 100644 index 0000000..2025047 --- /dev/null +++ b/miles/modal_train_glm_test.py @@ -0,0 +1,144 @@ +"""Dedicated 5-layer GLM-5.2 LoRA smoke-test launcher. + +glm5_2_744b_a40b_lora_5layer run on 1x8h200 + +Run: + uv run modal run miles/modal_train_glm_test.py::download_model + uv run modal run miles/modal_train_glm_test.py::download_data + uv run modal run -d miles/modal_train_glm_test.py::train +""" + +import asyncio +import os +import tempfile + +import modal + +from configs import get_module +from configs.base import HF_CACHE_PATH, DATA_PATH, CHECKPOINTS_PATH + +_ALLOWED_EXPERIMENTS = { + "glm5_2_744b_a40b_lora_5layer", +} +EXPERIMENT = os.environ.get("EXPERIMENT_CONFIG", "glm5_2_744b_a40b_lora_5layer") +if EXPERIMENT not in _ALLOWED_EXPERIMENTS: + raise ValueError(f"This launcher only supports 5-layer GLM diagnostics; got {EXPERIMENT!r}") + +exp_mod = get_module(EXPERIMENT) +modal_cfg = exp_mod.modal +miles_cfg = exp_mod.miles + +MILES_ROOT = "/root/miles" + +image = ( + modal.Image.from_registry(modal_cfg.docker_image) + .entrypoint([]) + .add_local_python_source("configs", copy=True) + .add_local_python_source("modal_helpers", copy=True) +) +for patch in modal_cfg.patch_files: + image = image.add_local_file( + patch, f"/tmp/{os.path.basename(patch)}", copy=True + ) +if modal_cfg.image_run_commands: + image = image.run_commands(*modal_cfg.image_run_commands) +if modal_cfg.image_env: + image = image.env(modal_cfg.image_env) + +with image.imports(): + from ray.job_submission import JobSubmissionClient + from modal_helpers.utils import ( + build_train_cmd, + prepare_miles_config, + start_ray_head, + ) + +hf_cache_volume = modal.Volume.from_name("huggingface-cache", create_if_missing=True) +data_volume = modal.Volume.from_name("miles-data", create_if_missing=True) +checkpoints_volume = modal.Volume.from_name("miles-checkpoints", create_if_missing=True) + +modal_volumes = { + str(HF_CACHE_PATH): hf_cache_volume, + str(DATA_PATH): data_volume, + str(CHECKPOINTS_PATH): checkpoints_volume, +} + +app = modal.App(f"{EXPERIMENT}-test") + +RAY_DASHBOARD_PORT = 8265 + + +def run_config_hook(experiment: str, hook_name: str, mounted_volumes) -> None: + cfg = get_module(experiment).miles + for volume in mounted_volumes: + volume.reload() + getattr(cfg, hook_name)() + for volume in mounted_volumes: + volume.commit() + + +@app.function( + image=image, + volumes={str(HF_CACHE_PATH): hf_cache_volume}, + timeout=4 * 60 * 60, + secrets=[modal.Secret.from_name("huggingface-secret")], +) +def download_model(experiment: str = EXPERIMENT): + run_config_hook(experiment, "download_model", (hf_cache_volume,)) + + +@app.function( + image=image, + volumes={str(DATA_PATH): data_volume}, + timeout=4 * 60 * 60, + secrets=[modal.Secret.from_name("huggingface-secret")], +) +def download_data(experiment: str = EXPERIMENT): + run_config_hook(experiment, "download_data", (data_volume,)) + + +@app.function( + image=image, + gpu=f"{modal_cfg.gpu}:{miles_cfg.actor_num_gpus_per_node}", + memory=modal_cfg.memory if modal_cfg.memory else None, + cloud=modal_cfg.cloud if modal_cfg.cloud else None, + region=modal_cfg.region if modal_cfg.region else None, + volumes=modal_volumes, + secrets=[modal.Secret.from_name("wandb-secret")], + timeout=24 * 60 * 60, +) +async def train(experiment: str = EXPERIMENT): + await asyncio.gather( + hf_cache_volume.reload.aio(), + data_volume.reload.aio(), + checkpoints_volume.reload.aio(), + ) + exp_mod = get_module(experiment) + cfg = exp_mod.miles + my_ip = "127.0.0.1" + os.environ["MILES_HOST_IP"] = my_ip + os.environ["SGLANG_HOST_IP"] = my_ip + os.environ["HOST_IP"] = my_ip + + start_ray_head(my_ip, 1) + prepare_miles_config(cfg, tempfile.mkdtemp()) + + cmd = build_train_cmd(cfg, MILES_ROOT) + runtime_env = { + "env_vars": { + "no_proxy": f"127.0.0.1,{my_ip}", + "MASTER_ADDR": my_ip, + **cfg.environment, + } + } + + client = JobSubmissionClient("http://127.0.0.1:8265") + job_id = client.submit_job(entrypoint=cmd, runtime_env=runtime_env) + print(f"Job submitted: {job_id}") + print(f"Training {experiment} on 1 node x {exp_mod.modal.gpu}:{cfg.actor_num_gpus_per_node}") + print(f"Command: {cmd}") + + async with modal.forward(RAY_DASHBOARD_PORT) as tunnel: + print(f"Ray dashboard: {tunnel.url}") + async for line in client.tail_job_logs(job_id): + print(line, end="", flush=True) diff --git a/miles/sglang_fp8_lora_fix.py b/miles/sglang_fp8_lora_fix.py new file mode 100644 index 0000000..3e1c89f --- /dev/null +++ b/miles/sglang_fp8_lora_fix.py @@ -0,0 +1,92 @@ +"""Image-build patcher: fix LoRA-B buffer sizing on quantized column-parallel layers. + +Root cause of "LoRA B output dim ... does not match base partition prefix dim" +under --quantization fp8: + + * mem_pool.get_lora_b_shape derives the effective TP for non-MoE column + modules from _row_parallel_shard_tp, an INPUT-sharding probe + (input_size // input_size_per_partition). + * On bf16, UnquantizedLinearMethod never sets input_size_per_partition, so + the probe falls back to the global tp_size and the code path below then + corrects the output dim via the output-side probe. Works. + * Fp8LinearMethod.create_weights DOES set layer.input_size_per_partition + (== input_size for column-parallel layers, whose input is unsharded), so + the probe returns 1, the whole sharding branch is skipped, and LoRA-B is + sized at the FULL output dim while the base layer stays TP-sharded. + set_lora_info then fails (e.g. shared_experts.gate_up_proj: B=4096 vs + per-rank partitions [1024, 1024] -> 2048). + +Fix: for non-MoE column-parallel modules, use the base module's +output_size_per_partition (probed by _column_parallel_out_partition) as the +authoritative per-rank LoRA-B output dim. It exists on both bf16 and +quantized layers and is exactly what set_lora_info validates against. +Applied at image build via `python /tmp/sglang_fp8_lora_fix.py`. +""" + +from pathlib import Path + +P = Path("/sgl-workspace/sglang/python/sglang/srt/lora/mem_pool.py") + +OLD = ''' if ( + effective_tp_size > 1 + and module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES + and module_name not in REPLICATED_LINEAR_LORA_NAMES + ): + # If the base column-parallel module is fully REPLICATED (its actual + # output_size_per_partition still equals the full output_dim -- e.g. the + # dense MLP gate_up under --moe-dense-tp-size 1), its output is NOT + # sharded, so keep LoRA-B at the full output dim. Dividing by the global + # tp_size here undersizes B and crashes set_lora_info ("LoRA B output dim + # != base partition prefix dim"). Non-MoE only; MoE shards by moe_tp_size. + probed_out = ( + None + if self.is_moe_module(module_name) + else self._column_parallel_out_partition( + module_name, base_model, layer_idx + ) + ) + if probed_out is not None and probed_out == output_dim: + pass # replicated base: keep full B output dim + else: + output_dim = self._column_parallel_lora_b_per_rank_dim( + module_name, output_dim, effective_tp_size + ) +''' + +NEW = ''' if ( + module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES + and module_name not in REPLICATED_LINEAR_LORA_NAMES + and not self.is_moe_module(module_name) + ): + # The base module's output_size_per_partition is the ground truth + # for LoRA-B's per-rank output dim (replicated OR TP-sharded), and + # it is quant-independent. Do NOT gate this on the input-sharding + # probe above: Fp8LinearMethod sets input_size_per_partition == + # input_size on column-parallel layers, which makes that probe + # return 1 and previously skipped sharding entirely, sizing LoRA-B + # at the full output dim against a TP-sharded base and crashing + # set_lora_info ("LoRA B output dim != base partition prefix dim"). + probed_out = self._column_parallel_out_partition( + module_name, base_model, layer_idx + ) + if probed_out is not None: + output_dim = probed_out + elif effective_tp_size > 1: + output_dim = self._column_parallel_lora_b_per_rank_dim( + module_name, output_dim, effective_tp_size + ) + elif ( + effective_tp_size > 1 + and module_name not in ROW_PARALLELISM_LINEAR_LORA_NAMES + and module_name not in REPLICATED_LINEAR_LORA_NAMES + ): + # MoE modules keep the moe_tp_size sharding path. + output_dim = self._column_parallel_lora_b_per_rank_dim( + module_name, output_dim, effective_tp_size + ) +''' + +src = P.read_text() +assert src.count(OLD) == 1, f"expected 1 match, got {src.count(OLD)}" +P.write_text(src.replace(OLD, NEW)) +print("patched", P) diff --git a/miles/sglang_tp1_shared_expert_fix.py b/miles/sglang_tp1_shared_expert_fix.py new file mode 100644 index 0000000..5e60e9b --- /dev/null +++ b/miles/sglang_tp1_shared_expert_fix.py @@ -0,0 +1,72 @@ +"""Fix SGLANG_SHARED_EXPERT_TP1 double-add in deepseek_v2.py. + +With a TP1-replicated shared expert, upstream adds its output to the MoE +result "after the all-reduce" so each TP rank contributes it only once. +But `should_skip_post_experts_all_reduce` can skip that explicit all-reduce +in favor of a *deferred/replaced* reduction (FlashInfer AllReduce Fusion, +dp-attention's reduce-scatterv), in which case the "post-all-reduce" add +actually lands *before* the real reduction and the replicated shared output +is summed once per TP rank (8-32x here) — corrupting every MoE layer. + +Fix: fold the shared output into the pre-reduction add scaled by +1/tp_size. Any linear reduction (all-reduce, fused all-reduce, +reduce-scatterv) then reconstitutes exactly one copy. tp_size is a power +of two, so the bf16 scaling is exact and there is no precision cost. + +Verified empirically on GLM-5.2-FP8: TP1 + AllReduce Fusion and +TP1 + dp-attention both produced garbage output before this patch. +""" + +from pathlib import Path + +TARGET = Path("/sgl-workspace/sglang/python/sglang/srt/models/deepseek_v2.py") + +SENTINEL = "MILES_TP1_SHARED_EXPERT_FIX" + +# Pre-reduction arg: pass the scaled shared output instead of None when TP1. +OLD_FUSE_ARG = "None if self._shared_expert_tp1 else shared_output," +NEW_FUSE_ARG = ( + "(shared_output * (1.0 / self.tp_size) if shared_output is not None else None) " + "if self._shared_expert_tp1 else shared_output, # " + SENTINEL +) + +# Post-reduction adds: now double-counting (the scaled copy is already in), +# so neuter them. +OLD_POST_ADD_DUAL = ( + " if self._shared_expert_tp1:\n" + " final_hidden_states += shared_output\n" +) +OLD_POST_ADD_NORMAL = ( + " if shared_output is not None and self._shared_expert_tp1:\n" + " final_hidden_states += shared_output\n" +) +NEW_POST_ADD = ( + " if False: # " + SENTINEL + ": folded into pre-reduction add\n" + " final_hidden_states += shared_output\n" +) + + +def main() -> None: + src = TARGET.read_text() + if SENTINEL in src: + print("already patched, skipping") + return + + n_fuse = src.count(OLD_FUSE_ARG) + assert n_fuse == 2, f"expected 2 maybe_fuse TP1 args, found {n_fuse}" + src = src.replace(OLD_FUSE_ARG, NEW_FUSE_ARG) + + n_dual = src.count(OLD_POST_ADD_DUAL) + assert n_dual == 1, f"expected 1 dual-stream post-add, found {n_dual}" + src = src.replace(OLD_POST_ADD_DUAL, NEW_POST_ADD) + + n_normal = src.count(OLD_POST_ADD_NORMAL) + assert n_normal == 1, f"expected 1 forward_normal post-add, found {n_normal}" + src = src.replace(OLD_POST_ADD_NORMAL, NEW_POST_ADD) + + TARGET.write_text(src) + print(f"patched {TARGET}: 2 pre-reduction args, 2 post-adds neutered") + + +if __name__ == "__main__": + main()