[Lilo][Miles runtime] Multi-node LoRA support and 256k Qwen3.8-27B definitions by micahtyong · Pull Request #39 · modal-projects/lilo

micahtyong · 2026-09-18T03:37:27Z

Summary

Adds the Qwen3.8-27B 256k definition and the multi-node trainer support it needs in the Miles LoRA backend. 256k on a 27B model doesn't fit on one node. qwen3_8_27b_miles_lora_256k is a 2-node TP2×CP8×DP1 (16×H200) trainer that mirrors the raw-Miles baseline exactly, so the two compare on identical hardware; rollout is 2×H200 TP4.

Validated e2e and benchmarked against another open source training framework (report)

Multi-node checkpoints are the substance of this PR

Definitions set TRAINER_NODES and run under @modal.experimental.clustered(N, rdma=True). Rank 0 becomes the Ray head and the Lilo engine; the other nodes ray start --address=head and idle until the head exits.

A Modal Volume has one view per container, and every part of the checkpoint path got that wrong in a different way. The first 256k run (lilo-27b-256k-pad-30-1789719362) trained 18 steps and then could never resume: checkpoint 000018 held only __0_0…__7_0, and every one of the nine chained resumes died with FileNotFoundError: /checkpoints/000018/.../miles/__16_0.distcp. Four distinct bugs were behind that, each fixed with a regression test:

Representative election by hostname selects nothing. _sync_checkpoint_volume() commits once per node, from one rank of each, elected by socket.gethostname() — but every container in a Modal cluster reports "modal", so hosts.index(...) == rank only ever picks global rank 0 and nodes 1..N-1 never commit. _node_identity() now elects by MODAL_TASK_ID (distinct per container, inherited by the Ray workers it spawns), falling back to the hostname off Modal.
Rank 0's rename publishes only rank 0's node. Miles has every rank write into _tmp_<ckpt>/ and rank 0 rename it into place; that rename runs against rank 0's local view, which contains 8 of 16 shards. Publication now goes through the volume's committed state:
```
write_shards(tmp); _sync_checkpoint_volume("commit")      # every node commits its shards
on rank 0:  volume.copy_files(peer_shards, final, recursive=True)   # server-side, no local view needed
            shutil.copy2(own_shards, final)                         # keeps rank 0's own view complete
```
The local copy of rank 0's own files matters: bulletin.publish reads the sampler adapter back locally right after writing it, and a server-side copy makes a file visible in committed state but not in the writing container.
Committing must not reload. A post-commit volume.reload() can never succeed on the node whose colocated engine holds a capture file open (there are open files preventing the operation: … __15_0.distcp is open from lilo.engine.backend_http), which killed a save mid-flight and surfaced client-side as backend /persist_snapshot transport failed (RemoteProtocolError). Writers don't need their peers' shards, so commit no longer refreshes the committing node; readers still reload on the load path.
/persist_snapshot never went through any of the above. MilesCommandBackend.persist_checkpoint() installs the capture with a local shutil.copytree, bypassing the Miles write_checkpoint_dir patch entirely — so even with 1–3 fixed, every checkpoint the running job wrote was still node-0-only. It now installs from committed volume state, and refuses the save outright when the capture holds fewer shards than the world size:
```
if shards and len(shards) < world_size:
    raise RuntimeError(f"capture {src} holds {len(shards)} of {world_size} shards; ...")
```
A checkpoint short of a node's shards loads on no rank but the ones that wrote it, and without this guard it is discovered only hours later, on the resume that needs it. _publish_checkpoints_across_nodes() carries the same check for the shard-write path, and rank-0 publication failures are broadcast so peers fail instead of hanging until the distributed timeout.

Trainer actors also had to be given LILO_CHECKPOINT_VOLUME / LILO_BULLETIN_VOLUME explicitly (_WORKER_ENV_VARS) — Ray workers on the non-head nodes didn't inherit them, so they silently did no volume ops at all. Before placing actors, _require_cluster_nodes polls ray.nodes() until actor_num_nodes alive GPU nodes expose world_size GPUs (Ray registers worker resources asynchronously), so a missing node fails fast with a clear message instead of hanging in placement.

256k-only trainer knobs (opt-in; smaller rungs are unchanged)

--distributed-timeout-minutes 120 — one 250k-token forward/backward microbatch takes longer than NCCL's 10 min default, which is what killed the first attempt.
Engine(max_forward_backward_batch=1) — caps how many compatible forward_backward requests the scheduler coalesces into one backend call; None keeps the old unbounded behaviour.
align_sequences_to_parallel_layout — pads each sequence to a multiple of 2*cp*tp so Megatron's zigzag CP chunks tile the TP ranks exactly, then trims the outputs back. Needed at TP8×CP3 (Miles only pads per-sequence to 2*cp); TP2×CP8 divides evenly so no shipped definition sets it. Kept as the workaround until the padding fix lands upstream in Miles.

Evidence

2-node TP2×CP8×DP1 fits 256k with headroom: per-rank max_allocated 82.4 GB / max_reserved 118.3 GB of 141 GB. A third node was never a memory requirement.
Matched 20-step run lilo-27b-256k-2n-20step-1790044500, 20/20 steps: mean cmp/train_time_s 2424.7 s (2403.7 s excluding step 0) vs raw Miles yz1wuem8 ~2125 s steady state — ~13% slower, not the ~77% reported when Lilo ran TP8×CP3 on 24 GPUs. End-to-end step time is faster (2448 s vs 2776 s) because Lilo's non-train time is ~44 s against Miles' ~599 s. Reward 0.485 mean, 814 tokens/GPU/s, prompt_len 250.4k, truncated_ratio 0. ~14 h wall, ~341 H200-hours.
The residual gap is Lilo's 128 separate forward_backward calls: 128 × 17.4 s + ~180 s assemble/publish = 2404 s. Left as-is deliberately — coalescing would change what's being compared.
Post-fix, a resume loaded a 16-shard multi-node checkpoint (load_slot ok=true) and subsequent saves write all 16 .distcp shards, which is exactly the failure the original run died on. The 20-step matched run lilo-27b-256k-2n-20step-1790044500 on this branch wrote complete 16-shard checkpoints at every one of its 20 saves (000000…000019, ~2.1 GiB each), step 18 included.

The control-plane fix for reaping wedged engines is split out into #53 — it was an unrelated bugfix that this work happened to surface.

Link to Devin session: https://modal.devinenterprise.com/sessions/ca503f6afb54412498a66f69af517adf
Open in Devin Desktop: https://modal.devinenterprise.com/desktop/session/ca503f6afb54412498a66f69af517adf?variant=devin
Requested by: @micahtyong

devin-ai-integration · 2026-09-18T03:37:31Z

🤖 Devin AI Engineer

I'll be helping with this pull request! Here's what you should know:

✅ I will automatically:

Address comments on this PR that start with 'DevinAI' or '@devin'.
Look at CI failures and help fix them

Note: I can only respond to comments from users who have write access to this repository.

⚙️ Control Options:

Disable automatic comment, CI, and merge conflict monitoring

micahtyong · 2026-09-18T04:09:37Z

/devin review

devin-ai-integration · 2026-09-18T04:09:40Z

Starting Devin Review.