Repository navigation
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Claude Code Review
This repository is configured for manual code reviews. Comment @claude review for a one-time review, or @claude review always to subscribe this PR to a review on every future push.
Tip: disable this comment in your organization's Code Review settings.
HJSang
changed the title
5aee0fa to
55ba5dd
Compare
HJSang
changed the title
37f99f2 to
6e10e42
Compare
73ccc35 to
792c4be
Compare
792c4be to
dbbf24a
Compare
|
@claude review |
dbbf24a to
9c399bb
Compare
HJSang
changed the title
8820713 to
f64aa02
Compare
f64aa02 to
a43f1ac
Compare
a43f1ac to
4144bf7
Compare
ff35dff to
63d3daa
Compare
|
@claude review |
|
@codex review |
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
|
Codex Review: Didn't find any major issues. Chef's kiss. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
There was a problem hiding this comment.
Findings marked 🟡 are optional suggestions and need no follow-up push.
Still open from earlier reviews (1):
- Unresolved: 1 minor or pre-existing.
Additional findings (outside the current diff — GitHub can't attach inline comments there):
-
🔴
miles/backends/training_utils/weight_update/protocols/shm.py— On IPv6 clusters, shm mode fails at connect even when every engine shares the trainer's host._engines_off_this_hostresolves each engine host withsocket.gethostbyname, which supports only IPv4. An IPv6 literal such as[fd00::1](fromwrap_ipv6) raisessocket.gaierrorinstead of returning False._is_local_addressalso probes with an AF_INETsocket.socket(), so an IPv6 address can never bind. Fix: resolve withsocket.getaddrinfoand probe with a socket of the matching family, so IPv4, IPv6 and dual-stack hosts all classify correctly.Why this was flagged
The trigger is
--update-weight-transfer-mode shmon a cluster whose engine URLs carry IPv6 addresses. The repo builds those withwrap_ipv6(kubernetes/core/cell_view.py:132, ray_worker_manager.py:364, external_engine_provider.py:197). On every rank,connect(shm.py:57) calls_engines_off_this_host. That function passesurlparse(server_url).hostnametosocket.gethostbyname(shm.py:124), which raisessocket.gaierrorfor IPv6 addresses. The exception escapesconnect_rollout_engines, so the trainer fails at the first connect or reconnect, even when the engines really are on the trainer's host. Even past name resolution,_is_local_address(shm.py:358) creates an AF_INET socket, so binding an IPv6 address always fails and the engine would be reported as "run elsewhere". The base branch has no shm mode; the other modes do not do this check, so they work on the same IPv6 cluster.Verification: Trigger:
--update-weight-transfer-mode shmwith IPv6 engine URLs. In miles/backends/training_utils/weight_update/protocols/shm.py, line 124 callssocket.gethostbyname, which raisessocket.gaierrorforfd00::1._is_local_address(line 130) probes with AF_INETsocket.socket(), so IPv6 never binds. The base branch has no shm mode, so this is a new mode that does not work in a supported setup. -
🟡
miles/backends/training_utils/weight_update/protocols/utils/data_replica.py— On CUDA trainers, disk and shm syncs pin model-sized host memory and never release it between syncs. Thetake()docstring promises nothing stays alive after the sync, but this memory does.tensor.to("cpu", non_blocking=True)on a CUDA tensor allocates pinned memory through PyTorch's caching host allocator. That allocator keeps freed blocks cached instead of returning them to the OS. Fix: copy into pageable host memory, or free the host allocator cache aftertake()'s buckets are consumed, so host RAM drops back once the sync ends. The CPU-only fast tests must keep passing.Why this was flagged
The trigger is a CUDA trainer with
--update-weight-transfer-mode diskorshm. Every rank receives every gathered bucket (is_sender = True, disk.py:68, shm.py:67).OwnedBuckets.offer(data_replica.py:41) copies its owned share to the host withnon_blocking=True. For a CUDA source and a CPU destination, ATen's_to_copysets pin_out, so each copy is a page-locked block from the CachingHostAllocator, rounded up in size. In disk mode all owned buckets stay alive untilafter_base_weights. Across the ranks on a node that is up to one full model copy (with full gather,replica_size = world_size, data_replica.py:15). After the buckets are dropped, the allocator caches those pinned blocks for reuse and does not unpin them. So the trainer host permanently loses that RAM, and it cannot be swapped. The PR text says pinned memory was not wanted ("pinned memory does not change it on torch_tpu"). The base branch has no such path.Verification: Trigger: a CUDA trainer with
--update-weight-transfer-mode diskorshm.OwnedBuckets.offer(data_replica.py:41) runstensor.to("cpu", non_blocking=True), so each host copy is pinned memory from the CachingHostAllocator, which keeps freed blocks cached. Nothing in the diff empties it. That memory stays reserved in the trainer process between syncs.
|
Addressed the two findings outside the diff in a3e9d0c: IPv6 engine hosts now resolve with getaddrinfo and probe a socket of the matching family; owned buckets copy non-blocking into unpinned host tensors, so no pinned memory stays cached. Thanks. |
|
@codex review |
|
Codex Review: Didn't find any major issues. Nice work! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
…DA IPC An engine on another accelerator stack (sglang-torchtpu) shares neither NCCL nor CUDA IPC with the trainer, and serves no /pull_weights, so none of the existing transfer modes reach it. --update-weight-transfer-mode disk writes the gathered HF buckets as a fresh safetensors checkpoint on a shared filesystem and has the engines reload it through /update_weights_from_disk. --update-weight-transfer-mode shm copies each flattened bucket into a shared-memory tensor and hands its handle to engines on the trainer's hosts through /update_weights_from_tensor, so nothing is written or parsed. In both, the ranks holding identical gathered buckets split them round-robin and copy only after the collective gathers, so the host copies run in parallel. shm runs inside the begin/end weight-update session; disk pauses the engines around its one reload call, as disk-delta does. Engine calls go through the rollout-cell updaters, so a failed engine is reported per cell. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…owth - Copy each owned bucket to the host when it is offered instead of holding it on the device until the stream ends, which added model/replica_size to every rank's peak device memory. - Raise a rank's write, publish or reload error on every rank, so the others do not wait in a collective the failed rank never joins. - Clear the versions a crashed run left in --update-weight-disk-dir on the first sync; each is a full checkpoint. - Fail at connect when an shm engine is not on the trainer's host, and in validation when an shm engine spans nodes. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A worker must stay reachable over rpc, so it may not import ray (tests/fast/utils/workers/test_ray_communication_boundary.py). Binding the engine's address answers the actual question, whether it belongs to this machine, also on hosts with several addresses. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…c for multi-policy
- Resolve shm engine hosts with getaddrinfo and probe each address with a
socket of its family, so IPv6 engine URLs no longer fail at connect.
- Copy owned buckets non-blocking into unpinned host tensors:
`.to("cpu", non_blocking=True)` pins memory that CUDA's host allocator keeps
after the sync. A blocking copy would also avoid it, but took Qwen3-4B shm
syncs on TPU from 3.1-3.9 s to 6.5-6.7 s.
- Refuse --update-weight-transfer-mode disk in multi policy runs, where every
policy would write the same version directories.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Each bucket goes to the replica rank with the fewest bytes so far, as Megatron's fully parallel save and PyTorch DCP's dedup_save_plans balance replicated writes; round-robin by count let an embedding-sized bucket skew one rank's share. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
6e84a57 to
3daf866
Compare
|
@claude review |
|
@codex review |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3daf86628d
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # a crashed run leaves full checkpoints behind, and partial shards in the version a resume rewrites | ||
| for stale_dir in glob.glob(os.path.join(self.root_dir, "weight_v*")): | ||
| shutil.rmtree(stale_dir, ignore_errors=True) | ||
| os.makedirs(self._version_dir, exist_ok=True) |
There was a problem hiding this comment.
In a multi-rank run, if the shared filesystem is unwritable, full, or unavailable, os.makedirs raises only on rank 0. Other ranks return from begin_sync and enter _mark_cells_errored_on_any_rank in WeightUpdater.update_weights, whose collective rank 0 will never join, leaving the workers blocked until the distributed timeout instead of reporting the filesystem failure promptly. Wrap the rank-0 directory preparation in fail_on_every_rank, as the subsequent shard-writing and reload steps already do.
Useful? React with 👍 / 👎.
| updater.submit_unbounded_client_call( | ||
| "update_weights_from_disk", model_path=self._version_dir, weight_version=str(weight_version) | ||
| ) |
There was a problem hiding this comment.
🔴 Disk-mode runs whose engines start with --sglang-load-format dummy (or MILES_SGLANG_DUMMY_LOAD=1) fail every weight sync, so the engines never receive the trained weights. The update_weights_from_disk call at disk.py:116-118 passes no load_format. When the request leaves it unset, SGLang falls back to the engine's startup load format, and its reload accepts only DefaultModelLoader. Fix: pass the format of the files disk mode writes, e.g. load_format="safetensors", so the reload never inherits the engine's startup loader.
Why this was flagged
This needs an SGLang build whose /update_weights_from_disk falls back to server_args.load_format and rejects loaders other than DefaultModelLoader. The repo starts engines that way when they get their weights from the trainer: sglang_engine.py:129-130 (MILES_SGLANG_DUMMY_LOAD), examples/infra_features/p2p_weight_transfer/run.py:984, and scripts/run_deepseek_v4.py:717. The trigger is --update-weight-transfer-mode disk with dummy-loaded engines. _reload_engines (disk.py:110-122) then gets HTTP 400 from every engine, and _make_request raises. With FT off, mark_errored re-raises and every rank fails at the first sync, leaving the engines paused. With weight_update FT on, every cell is marked errored and keeps serving dummy weights. No argument check rejects this pairing. On the base branch disk mode does not exist, and broadcast/shm push the tensors directly, so they work with dummy-loaded engines.
Verification: _reload_engines (disk.py:116-118) passes no load_format. SGLangApiClient.update_weights_from_disk (sglang_api_client.py:344-358) adds load_format to the payload only when it is not None. sglang_engine.py:129-130 sets kwargs["load_format"] = "dummy" when MILES_SGLANG_DUMMY_LOAD == "1". The new disk asserts in arguments.py:1296-1310 do not cover the load format.
Why
An engine on another accelerator stack, such as sglang-torchtpu, shares neither NCCL nor CUDA IPC with the trainer and serves no
/pull_weights, so no existing--update-weight-transfer-modecan reach it. This adds two modes built on endpoints every sglang engine already serves. The TPU stack above (#4213 onward) syncs through them.What changes
Two new values of
--update-weight-transfer-mode; existing modes and the default (broadcast) are unchanged.disk--update-weight-disk-dir; the engines reload it/update_weights_from_diskshm/dev/shmtensor that the engines map in place/update_weights_from_tensor(flattened_bucket)Reading order:
protocols/utils/data_replica.py(OwnedBuckets): after each gather, every rank of a pipeline stage holds the same bucket; each bucket goes to the replica rank with the fewest bytes so far (the rule Megatron's fully parallel save and PyTorch DCP'sdedup_save_plansuse), so the ranks share the host copies and writes. A rank copies its buckets to unpinned host tensors as they arrive (non-blocking), so it does not hold its share of the model on the device until the stream ends.protocols/disk.py: each rank writes its shards; rank 0 merges the index, copies the non-weight files of--hf-checkpoint, pauses the engines, reloads them and resumes, then deletes the previous version. The first sync also clears versions a crashed run left. One call replaces every weight, so it runs outside the begin/end session and without the per-request timeout (submit_unbounded_client_call).protocols/shm.py: runs inside the regular weight-update session. Each bucket is laid out on the host in the engine's flattened layout (from meta tensors; flattening on the device fails on TPU, which pads each tensor's last dimension). At most a few buckets are in shared memory at once. Each receiving process gets its own serialization, since each takes its own reference to the/dev/shmfile.connectfails fast when an engine is not on this host.protocols/utils/rank_errors.py(fail_on_every_rank): a write, publish or reload error on one rank is raised on every rank, so none waits in a collective the failed rank never joins (the patterndisk-deltauses inline).protocol.py,arguments.py,multi_policy/utils.py: selection and startup checks.broadcast.py's flattening helper moves toprotocols/utils/flattened_bucket.py, shared withshm.Startup checks
--colocate, no LoRA, no PD disaggregationshm: one trainer node, each engine on one node,--sglang-pp-size 1/dev/shmname; PP stages reuse TP ranks, so one payload would be opened twicedisk:--update-weight-disk-dirset,--hf-checkpointa local dir, no speculative draft unless--enable-mtp-trainingdisk: not in multi-policy runsTests
Weights arrive exactly: with
--check-weight-update-equal, each engine snapshots the checkpoint weights, overwrites them, and the first sync must restore every tensor bit for bit.diskshmbroadcast(reference)The TPU check needs sglang-torchtpu's weight checker to compare off CUDA (fix pending upstream); the sync itself needs nothing new from the engine.
Sync time and memory on TPU v7x, Qwen3-4B:
shmsyncdisksync,/dev/shmtarget / NFS/worktargetshm3.1–4.1 s vs 9.8–10.1 s;diskon NFS about equalUnit tests:
test_disk.py(each tensor written once and found through the merged index, reload between pause and resume, previous and stale versions deleted, failed reload raises),test_shm.py(another process rebuilds every TP rank's bucket, each receiver can open the file after the others released it, the window bounds host copies, the sharing strategy is restored, an errored cell gets no payload, only engines on this host pass),test_data_replica.py(ownership balances bytes),test_rank_errors.py,test_rollout_cell_updater.py(the unbounded call outlives the per-request timeout), and the multi-policy refusal.Limits and follow-ups
shmreaches only engines on the trainer's host; cross-host sync goes throughdisk. A network or ICI transport is a follow-up.🤖 Generated with Claude Code