Repository navigation
p2p: derive HF-to-parameter mappings from the model loader - #3969
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.
yueming-yuan
changed the title
Its only user, miles' p2p sender, now finds where each HF tensor lands by running every replica's own load_weights (radixark/miles#3969).
yueming-yuan
added
run-ci-weight-update
yueming-yuan
changed the title
ParameterMapper copied each model's loader renames and fusions by hand and exists for 15 of 244 model files. A name it did not know mapped to itself, the stager dropped it with a warning and the update passed: no MTP draft class has a mapper that matches its loader (Qwen3.5 MTP, DeepSeek and GLM-lite NextN, GLM4 NextN), and neither does the Qwen3.5 target. Each replica now runs its own load_weights once over tensors of every HF name the trainer sends, made one at a time on the trainer's GPU, its params swapped for meta ones in their reload layouts, and records which params each name's data is written into (traced through views, cat and dequant; padding writes carry no name). The names come from one pass of the real HF iterator at the first sync. The stager waits for every name of a param group (params sharing HF tensors), packing keeps a group in one buffer, the transfer buffers fit the largest group, and a sync sending a name outside the first pass fails. Buffers the loader fills besides params, such as the weight + 1 of sglang's Gemma norm, get scratch storage on the transfer buffer device; the engine recomputes them after an update.
yueming-yuan
removed
run-ci-weight-update
b8f4225 to
e89a2c9
Compare
yueming-yuan
changed the title
Keep aliasing, collective and storage constraints near the code; leave the repeated pipeline explanation in the PR and documentation.
yueming-yuan
changed the title
There was a problem hiding this comment.
[P1] No check that every published param is actually written. HF names that no replica maps are logged at INFO and skipped (_log_hf_names_no_runner_loads), and nothing asserts the converse: that every param the engine publishes is reached by at least one HF name. ModelParamStager.assert_all_done cannot catch it either — a group only enters _missing_hf_names_by_param_group once one of its HF names is staged, so a group with zero arrivals is invisible. Any name the trainer's converter emits that the engine loader ignores (or a write the dispatch-mode recorder misses) leaves that weight at its initial value forever, with no error. e2e with --check-weight-update-equal would catch it; production would not. Suggest a per-role assert at first sync: published param names ⊆ hf_name_mapping.hf_names_by_param_name (aggregated across senders when PP > 1), with an explicit allowlist for params that are intentionally not updated.
[P2] Transfer buffers are sized only at the first begin_sync. After a reconnect that builds new replicas, map_hf_names merges new mappings that can form larger param groups, and pack_into_buffers then fails with "need X bytes, a buffer has Y". Either re-check/resize on new mappings or fail at connect with a clear message.
|
@claude review always |
There was a problem hiding this comment.
Nothing blocking. The comments below are optional suggestions. There is no need to push a fix for them before merging.
Beyond the inline finding, I also checked whether buffer sizing could hard-fail once a new shard layout connects after the first sync for elastic/PD rollout fleets — investigated and ruled out as not a new regression introduced here.
Extended reasoning...
This PR replaces the static HF-to-parameter mapping in P2P weight transfer with a runtime probe (loader_probe.py) and reworks model_replica.py/model_param_stager.py/p2p.py to stage and transfer by parameter group; it touches the hot sync path but no auth/crypto/exposure surface. A confirmed correctness/perf finding already flagged (begin_sync's unconditional materialize=True probe pass runs on every rank, not just senders) is being posted inline, so a human look is warranted; I additionally examined the buffer-resizing concern for elastic/PD fleets raised as a candidate and ruled it out as pre-existing/not a regression from this change.
| if self._hf_tensor_specs is None: | ||
| self._hf_tensor_specs = { | ||
| hf_name: HfTensorSpec(tuple(tensor.shape), tensor.dtype) | ||
| for bucket in iter_buckets(materialize=True) |
There was a problem hiding this comment.
🟡 (optional) Every training rank, not just P2P senders, now pays the full weight-conversion cost (dtype cast, quantization) once per training run, which the base never charged non-senders. begin_sync()'s probe pass at p2p.py:74 calls iter_buckets(materialize=True) unconditionally for all ranks (line 71-76 run before the if self.is_sender: check at line 77), whereas the normal send loop in updater.py always passes materialize=protocol.is_sender, so non-senders previously only joined collectives (materialize=False) and skipped per-tensor conversion. In hf_weight_iterator_bridge.py:79-83, materialize=False explicitly skips _postprocess_and_quantize; forcing True there makes every non-sender rank run it once. Fix: gate the probe's full materialization so non-sender ranks still join the collective but skip per-tensor postprocess/quantize, e.g. build specs from a lightweight shape/dtype pass instead of materialize=True for every rank.
Why this was flagged
Trigger: the first update_weights() call of a training run (weight_version 0->1) on a multi-rank P2P job where some ranks are non-senders (assign_rollout_engine_ranks leaves assignments empty for them). begin_sync (p2p.py:67-76) runs on every rank and always passes materialize=True to iter_buckets, regardless of is_sender. In miles/backends/megatron_utils/update_weight/hf_weight_iterator_bridge.py:79-83, materialize gates whether _postprocess_and_quantize (expensive per-tensor dtype/quantization work) runs at all; with materialize=False (the base's path for non-senders) it is skipped. Now every non-sender additionally runs it once, during the first sync only, costing extra GPU memory/time for ranks that previously only joined cheap collectives. No existing guard limits this to senders.
Verification: nit. p2p.py:69-76 begin_sync runs on every rank (called unconditionally at updater.py:127) and, before the if self.is_sender: guard at p2p.py:77, iterates iter_buckets(materialize=True). In hf_weight_iterator_bridge.py:79-85, materialize=False only joins TP collectives, while materialize=True also runs _postprocess_and_quantize.
# Conflicts: # miles/backends/training_utils/weight_update/protocols/p2p.py # miles/backends/training_utils/weight_update/protocols/utils/model_replica.py
begin_sync now probes the iterator's buckets, and a param's first load runs twice.
There was a problem hiding this comment.
No new blocking issues. 1 optional suggestion (a nit or a note on pre-existing code) was found and not posted. Nothing in this review needs a push before merging.
Still open from earlier reviews (1):
- Unresolved: 1 minor or pre-existing.
This review covers commit b056fc6, which is no longer the latest commit on this pull request; later commits are not covered by it.
# Conflicts: # docs/advanced/p2p-weight-transfer.md
# Conflicts: # miles/backends/training_utils/weight_update/protocols/utils/model_replica.py
Buffers were sized once, at the first sync, to the largest param group. A reconnect that builds replicas for a new layout merges their mappings into the existing ones, and two params that now share an HF tensor form one larger group; it no longer fit a buffer, so every later update failed in pack_into_buffers. begin_sync now recomputes the size the groups need and, when it grew, waits for the old buffers' writes, unregisters them from Mooncake and allocates new ones.
There was a problem hiding this comment.
Code review completed
Nothing new to post: everything this review found is already covered by existing comments on this pull request or didn't merit a separate one.
Still open from earlier reviews (1):
- Unresolved: 1 minor or pre-existing.
This review covers commit 5c48919, which is no longer the latest commit on this pull request; later commits are not covered by it.
|
Re: [P1] no check that every published param is written, in #3969 (review) We check this with We don't add a separate runtime check because getting, on the trainer side at runtime, the list of tensors the inference engine will accept is cumbersome. It needs a maintained list of exceptions (vision towers, the head of a tied-embedding model, MTP layers when no draft runner is written, ...), which is hard to keep right as models and formats change. (human reply) |
Replace the hand-written
ParameterMapperwith a probe of each replica'sload_weights. This follows the model's actual renames, fusions and expert routing instead of silently dropping weights missing from a manual map.Stacked on #3959; draft support follows in #3985. Requires sgl-project/sglang#42820, which removes
ParameterMapper.Recorded validation: 296 weight-update/API-client tests passed on H200 with sglang. Coverage includes loader dependencies, value-dependent operations, tied weights, complete-group staging and parameter restoration after errors. Cleanup passed
pre-commit run --all-files; GPU results and failure controls for the combined stack are in #3985.ci-sglang-pr: #42820