Adding GLM 5.2 LoRA to modal multinode by kailash109 · Pull Request #91 · modal-labs/multinode-training-guide

kailash109 · 2026-07-10T15:42:51Z

Added configs to replicate GLM 5.2 LoRA addition to Miles -- PR on radixark/miles#1559

8 x 8 H200 training config on GSM8K (TP8, EP32, DP8, PP/CP 1) and (untested) longer-context DAPO variant
Added 4-H200 smoke-test launcher on 5 layer toy model (done the same as in Miles PR)

Validation plots (experiment miles/configs/glm5_2_744b_a40b_lora.py)
Logprob diff:

Reward:

devin-ai-integration

Devin Review found 4 potential issues.

devin-ai-integration · 2026-07-10T15:49:23Z

+    dsa_attention_backend = "tilelang"
+    qkv_format = "thd"


🔴 Long-context training config uses a known NaN-producing attention backend, making training runs produce garbage gradients

The long-context DAPO training variant uses the tilelang+thd attention combination (dsa_attention_backend = "tilelang" and qkv_format = "thd" at miles/configs/glm5_2_744b_a40b_lora_dapo.py:44-45) which is explicitly documented as producing NaN gradients in the sibling config, so any training run with this config will diverge immediately.

Impact: Training runs launched with this config will produce NaN gradients and fail to learn anything.

The main GLM config documents the NaN issue and uses a workaround that this config missed

In miles/configs/glm5_2_744b_a40b_lora.py:58-61, the main config explicitly notes:

# 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"

The dapo config (miles/configs/glm5_2_744b_a40b_lora_dapo.py) uses the same full 744B model (hf_checkpoint = "zai-org/GLM-5.2") but sets dsa_attention_backend = "tilelang" and qkv_format = "thd" — the exact combination the main config avoids. The dapo config also omits data_pad_size_multiplier = 32 which the main config uses to control activation memory.

The 5-layer smoke-test config also uses tilelang+thd, but that is a toy model for path-testing, not production training.

Suggested change

dsa_attention_backend = "tilelang"

qkv_format = "thd"

dsa_attention_backend = "megatron"

qkv_format = "bshd"

data_pad_size_multiplier = 32

Was this helpful? React with 👍 or 👎 to provide feedback.

devin-ai-integration · 2026-07-10T15:49:26Z

+
+    actor_num_nodes = 8
+    actor_num_gpus_per_node = 8
+    num_gpus_per_node = 8       


🔍 Unique num_gpus_per_node field only appears in one config

This config defines num_gpus_per_node = 8 at line 70, which is not in _MILES_SKIP (miles/configs/base.py:31-37) and will be emitted as --num-gpus-per-node 8 via cli_args(). No other config in the entire repository (miles or slime) uses this field name — all others use actor_num_gpus_per_node. If --num-gpus-per-node is a valid Miles CLI flag, this is fine. If it's not recognized, Miles may reject it or silently ignore it. Worth confirming against the Miles CLI documentation.

Was this helpful? React with 👍 or 👎 to provide feedback.

devin-ai-integration · 2026-07-10T15:49:27Z

+    context_parallel_size = 1
+    expert_model_parallel_size = 8
+    expert_tensor_parallel_size = 1


🔍 The 8-node config uses moe_token_dispatcher_type = "alltoall" but the 1-node dapo config omits it

The main 8-node config (miles/configs/glm5_2_744b_a40b_lora.py:80) explicitly sets moe_token_dispatcher_type = "alltoall", which is important for MoE expert parallelism. The dapo config (miles/configs/glm5_2_744b_a40b_lora_dapo.py) and the 5-layer config both omit this field, presumably relying on the Miles default. If the default dispatcher type differs from alltoall, this could affect performance or correctness for MoE routing with EP > 1.

Was this helpful? React with 👍 or 👎 to provide feedback.

devin-ai-integration

Devin Review found 4 new potential issues.

devin-ai-integration · 2026-07-10T15:54:52Z

+    start_ray_head(my_ip, 1)
+    prepare_miles_config(cfg, tempfile.mkdtemp())
+
+    cmd = build_train_cmd(cfg, MILES_ROOT)


🔴 Weights & Biases API key not forwarded to the training job, so experiment logging fails

The wandb API key is never attached to the training config (modal_train_glm_test.py:116), so the generated CLI command omits the --wandb-key argument even though logging is enabled, causing the training job to fail to authenticate with Weights & Biases.

Impact: The smoke-test training run will either crash or silently skip all wandb experiment logging.

Missing wandb key forwarding compared to the main launcher

The main launcher at miles/modal_train.py:293-296 explicitly reads WANDB_API_KEY from the environment (injected by the wandb-secret Modal secret) and sets it on the config object:

if (wandb_key := os.environ.get("WANDB_API_KEY", "")) and getattr( miles_cfg, "use_wandb", False ): miles_cfg.wandb_key = wandb_key

This causes cli_args() to emit --wandb-key <key>. The test launcher at miles/modal_train_glm_test.py:107-116 skips this step entirely. Although the wandb-secret is mounted on the container (line 98), the Ray job's runtime_env at lines 117-122 only forwards no_proxy, MASTER_ADDR, and cfg.environment — none of which include WANDB_API_KEY. Ray jobs do not automatically inherit the parent process's full environment, so the key is lost.

The 5-layer config at miles/configs/glm5_2_744b_a40b_lora_5layer.py:134 sets use_wandb = True.

Suggested change

cmd = build_train_cmd(cfg, MILES_ROOT)

if (wandb_key := os.environ.get("WANDB_API_KEY", "")) and getattr(

cfg, "use_wandb", False

):

cfg.wandb_key = wandb_key

cmd = build_train_cmd(cfg, MILES_ROOT)

Was this helpful? React with 👍 or 👎 to provide feedback.

devin-ai-integration · 2026-07-10T15:54:54Z

+    dsa_attention_backend = "tilelang"
+    qkv_format = "thd"


🔍 DAPO config uses tilelang attention backend that the full 744B config explicitly avoids due to NaN gradients

The full 744B config at miles/configs/glm5_2_744b_a40b_lora.py:58-61 explicitly documents that tilelang + thd produces NaN gradients under nonzero loss, and switches to dsa_attention_backend = "megatron" with qkv_format = "bshd". However, the DAPO variant at miles/configs/glm5_2_744b_a40b_lora_dapo.py:44-45 uses dsa_attention_backend = "tilelang" and qkv_format = "thd" — the exact combination the full config warns against. The DAPO config docstring says "untested" and claims to be the "same recipe" as the full config, but this attention backend difference could cause NaN gradients during training. This may be intentional experimentation, but given the explicit warning in the sibling config, it deserves attention.

Was this helpful? React with 👍 or 👎 to provide feedback.

devin-ai-integration · 2026-07-10T15:54:55Z

+    sequence_parallel = True
+    pipeline_model_parallel_size = 1
+    context_parallel_size = 1
+    expert_model_parallel_size = 32


🔍 EP=32 exceeds standard Megatron DP constraint for 64-GPU topology

The full 744B config at miles/configs/glm5_2_744b_a40b_lora.py:78 sets expert_model_parallel_size = 32 with 64 total GPUs (8 nodes × 8 GPUs), TP=8, PP=1, CP=1. In standard Megatron-LM, DP = world_size / (TP × PP × CP) = 8, and EP must divide DP, so EP ≤ 8. EP=32 would be invalid under that constraint. However, the docstring at line 8 explicitly states "EP 32, DP 8" and says this was adapted from a validated launcher. Miles/Megatron-Bridge may use a different parallelism model where EP is not constrained by DP (e.g., EP could be orthogonal to DP). This is worth verifying against the Miles documentation.

Was this helpful? React with 👍 or 👎 to provide feedback.

kailash109 · 2026-07-16T15:50:51Z

Added FP8 rollout config as well, tested on dapo math:

logprob diff:

reward:

devin-ai-integration

Devin Review found 1 new potential issue.

devin-ai-integration · 2026-07-17T17:10:45Z

+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)."
+            )'''


🔍 megatron CP-assert gate depends on dsa_attention_backend being present on the transformer config

miles/megatron_dsa_cp_assert_fix.py:30-37 gates the DSA CP assert on getattr(self, "dsa_attention_backend", "megatron") == "tilelang". If the Megatron transformer config instance does not actually carry dsa_attention_backend at __post_init__ time, the getattr default ("megatron") keeps the assert enforced and CP>1 runs (16node/8node_cp4) would fail to start. This is fail-safe (errors out rather than corrupting), but the docstring's claim that finalize() propagates dsa_attention_backend onto the provider instance should be verified against the actual miles/megatron version in the pinned image.

Was this helpful? React with 👍 or 👎 to provide feedback.

kailash109 added 3 commits June 23, 2026 21:39

nemo rl multinode grpo example

94397d0

added working configs for glm5.2 lora on gsm8k

1dec971

Merge main into glm_5_2_lora

4565e7f

devin-ai-integration Bot reviewed Jul 10, 2026

View reviewed changes

revert modal_train

cab0aa5

devin-ai-integration Bot reviewed Jul 10, 2026

View reviewed changes

fp 8 rollout config

916b17c

This comment was marked as resolved.

Sign in to view

fp8 rollouts working on glm5.2 lora, misc bugfixing on sglang

56a6b3d

16 node wip config

f92ef7d

devin-ai-integration Bot reviewed Jul 17, 2026

View reviewed changes

-    dsa_attention_backend = "tilelang"
-    qkv_format = "thd"
+    dsa_attention_backend = "megatron"
+    qkv_format = "bshd"
+    data_pad_size_multiplier = 32

-    cmd = build_train_cmd(cfg, MILES_ROOT)
+    if (wandb_key := os.environ.get("WANDB_API_KEY", "")) and getattr(
+        cfg, "use_wandb", False
+    ):
+        cfg.wandb_key = wandb_key
+    cmd = build_train_cmd(cfg, MILES_ROOT)

Uh oh!

Conversation

kailash109 commented Jul 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

devin-ai-integration Bot left a comment

Choose a reason for hiding this comment

Uh oh!

devin-ai-integration Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

devin-ai-integration Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

devin-ai-integration Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

devin-ai-integration Bot left a comment

Choose a reason for hiding this comment

Uh oh!

devin-ai-integration Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

devin-ai-integration Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

devin-ai-integration Bot Jul 10, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

This comment was marked as resolved.

Uh oh!

kailash109 commented Jul 16, 2026

Uh oh!

devin-ai-integration Bot left a comment

Choose a reason for hiding this comment

Uh oh!

devin-ai-integration Bot Jul 17, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

kailash109 commented Jul 10, 2026 •

edited

Loading