refactor(model): unify Qwen + DeepSeek weight staging behind one declarative pipeline - #174
Conversation
|
Important
This repository does not receive automatic reviews because it has fewer than 10 stars. ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
…he new one hw-native-sys#163 step 8, which closes the issue. Net -1057 lines. **Deleted, now that the differentials that guarded them have run green in CI:** `_pack_deepseek_v4_layer_weights_legacy` and the five helpers only it used (`_replicate_weight`, `_pack_wo_a`, `_pack_deepseek_v4_optional_attention`, `_pack_deepseek_v4_router`, `_pack_deepseek_v4_routed_experts`), the three hand-written stacker helpers, and the two differential suites themselves. `weight_loader.py` goes from 1589 to 1097 lines. Keeping them longer was tempting and wrong: a differential compares two implementations, so the moment one is deleted the suite is comparing nothing. The proof is in the record — those tests ran green against the rule table on hw-native-sys#174 and on this branch — and `git log` has both sides if a byte-level question ever comes back. What remains is the property that outlives the refactor: the sidecar written by this writer round-trips through the reader and matches a fresh pack of the same checkpoint. **`LayerWeights` and `RuntimeModel.layers` are gone for real**, not just deprecated. The reason I gave for deferring this — that six open PRs construct a `RuntimeModel` — did not survive checking: none of them pass `layers=`. hw-native-sys#114's two matches are `num_layers=model.config.num_hidden_layers`, which is unrelated. Five call sites in tests and the loader dropped the argument. **Docs.** A new `docs/dev/model/weight-staging.md` covers the pipeline: the four stages, the six rule kinds, and the invariants that are not obvious from the code — rule order being sidecar contract, why nothing in the stacker calls `torch.cat`, why the stack axis differs per family, why DeepSeek stages serially and Qwen pooled, and why an LM head cannot pad with zeros. Both model docs gain a section pointing at it with their own specifics. The two traps found by running the tests rather than reading them are recorded there too, so they are not rediscovered: safetensors serializing its metadata in nondeterministic order (which makes a whole-file assertion flaky), and a synthetic checkpoint needing production dimensions for the compressor/indexer weights. 228 tests pass across `tests/unit`; `test_model_components.py` passes unmodified, which is the gate hw-native-sys#163 sets.
…he new one hw-native-sys#163 step 8, which closes the issue. Net -1057 lines. **Deleted, now that the differentials that guarded them have run green in CI:** `_pack_deepseek_v4_layer_weights_legacy` and the five helpers only it used (`_replicate_weight`, `_pack_wo_a`, `_pack_deepseek_v4_optional_attention`, `_pack_deepseek_v4_router`, `_pack_deepseek_v4_routed_experts`), the three hand-written stacker helpers, and the two differential suites themselves. `weight_loader.py` goes from 1589 to 1097 lines. Keeping them longer was tempting and wrong: a differential compares two implementations, so the moment one is deleted the suite is comparing nothing. The proof is in the record — those tests ran green against the rule table on hw-native-sys#174 and on this branch — and `git log` has both sides if a byte-level question ever comes back. What remains is the property that outlives the refactor: the sidecar written by this writer round-trips through the reader and matches a fresh pack of the same checkpoint. **`LayerWeights` and `RuntimeModel.layers` are gone for real**, not just deprecated. The reason I gave for deferring this — that six open PRs construct a `RuntimeModel` — did not survive checking: none of them pass `layers=`. hw-native-sys#114's two matches are `num_layers=model.config.num_hidden_layers`, which is unrelated. Five call sites in tests and the loader dropped the argument. **Docs.** A new `docs/dev/model/weight-staging.md` covers the pipeline: the four stages, the six rule kinds, and the invariants that are not obvious from the code — rule order being sidecar contract, why nothing in the stacker calls `torch.cat`, why the stack axis differs per family, why DeepSeek stages serially and Qwen pooled, and why an LM head cannot pad with zeros. Both model docs gain a section pointing at it with their own specifics. The two traps found by running the tests rather than reading them are recorded there too, so they are not rediscovered: safetensors serializing its metadata in nondeterministic order (which makes a whole-file assertion flaky), and a synthetic checkpoint needing production dimensions for the compressor/indexer weights. 228 tests pass across `tests/unit`; `test_model_components.py` passes unmodified, which is the gate hw-native-sys#163 sets.
58259ab to
a15e69d
Compare
…he new one hw-native-sys#163 step 8, which closes the issue. Net -1057 lines. **Deleted, now that the differentials that guarded them have run green in CI:** `_pack_deepseek_v4_layer_weights_legacy` and the five helpers only it used (`_replicate_weight`, `_pack_wo_a`, `_pack_deepseek_v4_optional_attention`, `_pack_deepseek_v4_router`, `_pack_deepseek_v4_routed_experts`), the three hand-written stacker helpers, and the two differential suites themselves. `weight_loader.py` goes from 1589 to 1097 lines. Keeping them longer was tempting and wrong: a differential compares two implementations, so the moment one is deleted the suite is comparing nothing. The proof is in the record — those tests ran green against the rule table on hw-native-sys#174 and on this branch — and `git log` has both sides if a byte-level question ever comes back. What remains is the property that outlives the refactor: the sidecar written by this writer round-trips through the reader and matches a fresh pack of the same checkpoint. **`LayerWeights` and `RuntimeModel.layers` are gone for real**, not just deprecated. The reason I gave for deferring this — that six open PRs construct a `RuntimeModel` — did not survive checking: none of them pass `layers=`. hw-native-sys#114's two matches are `num_layers=model.config.num_hidden_layers`, which is unrelated. Five call sites in tests and the loader dropped the argument. **Docs.** A new `docs/dev/model/weight-staging.md` covers the pipeline: the four stages, the six rule kinds, and the invariants that are not obvious from the code — rule order being sidecar contract, why nothing in the stacker calls `torch.cat`, why the stack axis differs per family, why DeepSeek stages serially and Qwen pooled, and why an LM head cannot pad with zeros. Both model docs gain a section pointing at it with their own specifics. The two traps found by running the tests rather than reading them are recorded there too, so they are not rediscovered: safetensors serializing its metadata in nondeterministic order (which makes a whole-file assertion flaky), and a synthetic checkpoint needing production dimensions for the compressor/indexer weights. 228 tests pass across `tests/unit`; `test_model_components.py` passes unmodified, which is the gate hw-native-sys#163 sets.
…he new one hw-native-sys#163 step 8, which closes the issue. Net -1057 lines. **Deleted, now that the differentials that guarded them have run green in CI:** `_pack_deepseek_v4_layer_weights_legacy` and the five helpers only it used (`_replicate_weight`, `_pack_wo_a`, `_pack_deepseek_v4_optional_attention`, `_pack_deepseek_v4_router`, `_pack_deepseek_v4_routed_experts`), the three hand-written stacker helpers, and the two differential suites themselves. `weight_loader.py` goes from 1589 to 1097 lines. Keeping them longer was tempting and wrong: a differential compares two implementations, so the moment one is deleted the suite is comparing nothing. The proof is in the record — those tests ran green against the rule table on hw-native-sys#174 and on this branch — and `git log` has both sides if a byte-level question ever comes back. What remains is the property that outlives the refactor: the sidecar written by this writer round-trips through the reader and matches a fresh pack of the same checkpoint. **`LayerWeights` and `RuntimeModel.layers` are gone for real**, not just deprecated. The reason I gave for deferring this — that six open PRs construct a `RuntimeModel` — did not survive checking: none of them pass `layers=`. hw-native-sys#114's two matches are `num_layers=model.config.num_hidden_layers`, which is unrelated. Five call sites in tests and the loader dropped the argument. **Docs.** A new `docs/dev/model/weight-staging.md` covers the pipeline: the four stages, the six rule kinds, and the invariants that are not obvious from the code — rule order being sidecar contract, why nothing in the stacker calls `torch.cat`, why the stack axis differs per family, why DeepSeek stages serially and Qwen pooled, and why an LM head cannot pad with zeros. Both model docs gain a section pointing at it with their own specifics. The two traps found by running the tests rather than reading them are recorded there too, so they are not rediscovered: safetensors serializing its metadata in nondeterministic order (which makes a whole-file assertion flaky), and a synthetic checkpoint needing production dimensions for the compressor/indexer weights. 228 tests pass across `tests/unit`; `test_model_components.py` passes unmodified, which is the gate hw-native-sys#163 sets.
9d9075d to
f166bb5
Compare
…he new one hw-native-sys#163 step 8, which closes the issue. Net -1057 lines. **Deleted, now that the differentials that guarded them have run green in CI:** `_pack_deepseek_v4_layer_weights_legacy` and the five helpers only it used (`_replicate_weight`, `_pack_wo_a`, `_pack_deepseek_v4_optional_attention`, `_pack_deepseek_v4_router`, `_pack_deepseek_v4_routed_experts`), the three hand-written stacker helpers, and the two differential suites themselves. `weight_loader.py` goes from 1589 to 1097 lines. Keeping them longer was tempting and wrong: a differential compares two implementations, so the moment one is deleted the suite is comparing nothing. The proof is in the record — those tests ran green against the rule table on hw-native-sys#174 and on this branch — and `git log` has both sides if a byte-level question ever comes back. What remains is the property that outlives the refactor: the sidecar written by this writer round-trips through the reader and matches a fresh pack of the same checkpoint. **`LayerWeights` and `RuntimeModel.layers` are gone for real**, not just deprecated. The reason I gave for deferring this — that six open PRs construct a `RuntimeModel` — did not survive checking: none of them pass `layers=`. hw-native-sys#114's two matches are `num_layers=model.config.num_hidden_layers`, which is unrelated. Five call sites in tests and the loader dropped the argument. **Docs.** A new `docs/dev/model/weight-staging.md` covers the pipeline: the four stages, the six rule kinds, and the invariants that are not obvious from the code — rule order being sidecar contract, why nothing in the stacker calls `torch.cat`, why the stack axis differs per family, why DeepSeek stages serially and Qwen pooled, and why an LM head cannot pad with zeros. Both model docs gain a section pointing at it with their own specifics. The two traps found by running the tests rather than reading them are recorded there too, so they are not rediscovered: safetensors serializing its metadata in nondeterministic order (which makes a whole-file assertion flaky), and a synthetic checkpoint needing production dimensions for the compressor/indexer weights. 228 tests pass across `tests/unit`; `test_model_components.py` passes unmodified, which is the gate hw-native-sys#163 sets.
|
|
||
| def _load_safetensors_dir(model_dir: Path) -> dict[str, torch.Tensor]: | ||
| """Load all safetensors shards from a local Hugging Face directory.""" | ||
| def _load_safetensors_subset( |
There was a problem hiding this comment.
We can make the ModelLoader a common base class for both qwen/deepseek to subclass it
There was a problem hiding this comment.
Good call — that is the same move this PR made for the stores (LazySafetensorsStore) and did not make for the loaders. Done in 33b3411.
SafetensorsDirectoryLoader now holds what is genuinely shared: supports_format was byte-identical in both, and both need a config.json before any of their own checks mean anything, so can_load becomes that precondition plus a _recognises hook per family.
I left load out rather than making it a template method, and it is worth saying why in case you disagree: the two produce different things from a directory — one stages Hugging Face layers lazily, the other validates a quantized checkpoint contract and returns metadata — so a shared skeleton would need a hook per step, which is more coupling than the duplication it removes. Happy to go further if you had a specific shape in mind.
One asymmetry kept deliberately and now commented: the DeepSeekV4 detection treats an unreadable config.json as "not mine" rather than an error, so the registry moves on to the next loader, while the Hugging Face one accepts any directory carrying safetensors.
| @@ -1197,393 +1053,45 @@ def pack_deepseek_v4_layer_weights( | |||
| destinations: Mapping[str, torch.Tensor] | None = None, | |||
| prefix: str | None = None, | |||
| ) -> DeepSeekV4PackedLayerWeights: | |||
There was a problem hiding this comment.
why we still kept this no-op func?
There was a problem hiding this comment.
Not a no-op, but the docstring made it look like one — it led with the history of the hand-written packer it replaced instead of what it does. Rewritten in 33b3411.
It is DeepSeekV4's binding to the generic evaluator: it turns the family's arguments into a LayerContext, selects the rank policy and the expert placement, supplies the synthetic-weight factories (the Hadamard index, the MTP smooth rows), and carries the two error templates that the existing diagnostics — and the tests matching on them — depend on. Both production callers (load_packed_layer_weights and load_mtp_weights) would otherwise repeat that wiring, and the MTP path would additionally have to know to thread prefix="mtp.0" through it.
If the concern is that the name now oversells it, I am happy to rename — but the argument surface is what the rest of the DeepSeekV4 code already calls, including tools/prepack_deepseek_v4.py, so I would rather not change the signature in this PR.
…n common Review on hw-native-sys#174 asked for a common base class for the Qwen and DeepSeekV4 loaders, which is the same move this PR made for the *stores* and had not made for the loaders. `SafetensorsDirectoryLoader` now holds the parts that are genuinely shared: `supports_format` was byte-identical in both, and both must find a `config.json` before any of their own checks mean anything — so `can_load` becomes that precondition plus a `_recognises` hook each family implements. `load` deliberately stays out of it. The two produce different things from a directory — one stages Hugging Face layers lazily, the other validates a quantized checkpoint contract and hands back metadata — and a shared skeleton would need a hook per step, which is more coupling than the duplication it removes. The DeepSeekV4 detection keeps its own shape rather than being folded into the base: it requires a shard index and a config that names DeepSeekV4, and treats an unreadable config as "not mine" rather than an error, so the registry moves on to the next loader. That asymmetry with the Hugging Face loader (which accepts any directory carrying safetensors) is intentional and now has a comment saying so. Also rewrote `pack_deepseek_v4_layer_weights`'s docstring, which the same review read as a no-op wrapper — fairly, because it led with the history of the hand-written packer it replaced instead of what it does. It is the family's binding to the generic evaluator: it builds the `LayerContext`, selects the rank and expert-placement policies, supplies the synthetic-weight factories, and carries the error templates its callers' diagnostics are matched against. Both production callers would otherwise repeat that wiring, and the MTP path would additionally have to thread `prefix="mtp.0"` through it. 257 tests pass; `ruff` and the repo hooks clean.
…n common Review on hw-native-sys#174 asked for a common base class for the Qwen and DeepSeekV4 loaders, which is the same move this PR made for the *stores* and had not made for the loaders. `SafetensorsDirectoryLoader` now holds the parts that are genuinely shared: `supports_format` was byte-identical in both, and both must find a `config.json` before any of their own checks mean anything — so `can_load` becomes that precondition plus a `_recognises` hook each family implements. `load` deliberately stays out of it. The two produce different things from a directory — one stages Hugging Face layers lazily, the other validates a quantized checkpoint contract and hands back metadata — and a shared skeleton would need a hook per step, which is more coupling than the duplication it removes. The DeepSeekV4 detection keeps its own shape rather than being folded into the base: it requires a shard index and a config that names DeepSeekV4, and treats an unreadable config as "not mine" rather than an error, so the registry moves on to the next loader. That asymmetry with the Hugging Face loader (which accepts any directory carrying safetensors) is intentional and now has a comment saying so. Also rewrote `pack_deepseek_v4_layer_weights`'s docstring, which the same review read as a no-op wrapper — fairly, because it led with the history of the hand-written packer it replaced instead of what it does. It is the family's binding to the generic evaluator: it builds the `LayerContext`, selects the rank and expert-placement policies, supplies the synthetic-weight factories, and carries the error templates its callers' diagnostics are matched against. Both production callers would otherwise repeat that wiring, and the MTP path would additionally have to thread `prefix="mtp.0"` through it. 257 tests pass; `ruff` and the repo hooks clean.
33b3411 to
6432384
Compare
…n common Review on hw-native-sys#174 asked for a common base class for the Qwen and DeepSeekV4 loaders, which is the same move this PR made for the *stores* and had not made for the loaders. `SafetensorsDirectoryLoader` now holds the parts that are genuinely shared: `supports_format` was byte-identical in both, and both must find a `config.json` before any of their own checks mean anything — so `can_load` becomes that precondition plus a `_recognises` hook each family implements. `load` deliberately stays out of it. The two produce different things from a directory — one stages Hugging Face layers lazily, the other validates a quantized checkpoint contract and hands back metadata — and a shared skeleton would need a hook per step, which is more coupling than the duplication it removes. The DeepSeekV4 detection keeps its own shape rather than being folded into the base: it requires a shard index and a config that names DeepSeekV4, and treats an unreadable config as "not mine" rather than an error, so the registry moves on to the next loader. That asymmetry with the Hugging Face loader (which accepts any directory carrying safetensors) is intentional and now has a comment saying so. Also rewrote `pack_deepseek_v4_layer_weights`'s docstring, which the same review read as a no-op wrapper — fairly, because it led with the history of the hand-written packer it replaced instead of what it does. It is the family's binding to the generic evaluator: it builds the `LayerContext`, selects the rank and expert-placement policies, supplies the synthetic-weight factories, and carries the error templates its callers' diagnostics are matched against. Both production callers would otherwise repeat that wiring, and the MTP path would additionally have to thread `prefix="mtp.0"` through it. 257 tests pass; `ruff` and the repo hooks clean.
6432384 to
c5ddc89
Compare
|
The DeepSeek accuracy cases (
Note this is a different failure from the intermittent prefix-cache assertion in #183; that one is not what is firing here.
|
|
Correcting my comment above: the attribution was wrong, and the issue I linked is closed as a duplicate. The What is still true: retrying this PR will not clear it, and it is not #183's flake. What I got wrong: it is no longer purely upstream. pypto's cap lives in I will put up the cap for our action separately; this PR needs it to go green but does not otherwise depend on it. |
…/weights First step of hw-native-sys#163, additive: the checkpoint reader is family-neutral, so it moves out of DeepSeek's weight loader and into a shared base the Qwen path can adopt in turn. `LazySafetensorsStore` owns the index (`filename_for`, `path_for`, `require`, `__contains__`) and the reads (`load_tensor`, `load_many`, grouped one open per shard); `DeepSeekV4WeightStore` subclasses it and keeps what is actually DeepSeek's — which names a checkpoint must expose, and how global and per-layer tensors are packed for the fused kernels. No behaviour change: 75 lines leave the loader, 21 arrive, and the loader's 1589-line pack path is untouched. Two seams were deliberate rather than incidental: - The three error messages are class attributes (`missing_name_error`, `missing_names_error`, `missing_shard_error`) instead of hard-coded strings, so DeepSeek keeps the diagnostics its users already recognise, word for word. The wording is part of a family's contract too, and a shared base that silently rewords it would be a regression nobody notices until they grep a log. - The default opener is resolved through `_default_open_fn()`, a method, not a module global read from the base. `DeepSeekV4WeightStore` overrides it to return this module's `_default_safe_open`, which carries the DeepSeekV4-specific import diagnostic and is what `test_model_components.py` monkeypatches by name. Reading the base module's global instead would have left that patch pointing at nothing and the test passing for the wrong reason. `tests/unit/model/common/test_weight_store.py` covers the generic store directly: shard grouping (interleaved names must not reopen a shard), caller order with duplicates dropped, the three message templates, a missing shard file, a real safetensors round trip, and a guard on the per-construction opener resolution above. Gate for this step, from hw-native-sys#163: `tests/unit/model/deepseek/test_model_components.py` passes unmodified — 81 tests, and 187 across `tests/unit`.
Step 1 of hw-native-sys#163, and deliberately first: the refactor has to keep the stacked slabs byte-identical, because pypto-lib fuses every layer into one kernel and treats the whole-model layout as an output contract. That is only enforceable if a test can fail on it, so the harness lands before any pipeline change. Three pieces: - `tests/unit/model/conftest.py` — `fingerprint_tensors`, a byte-level fingerprint (`name -> (shape, dtype, sha256)`). Hashing the bytes rather than comparing with `torch.equal` is what catches a dtype change that happens to round-trip through the same values. - `tests/unit/model/deepseek/conftest.py` — a synthetic on-disk checkpoint. Which tensors a layer needs comes from `deepseek_v4_layer_weight_names`, not from a copied list, so the fixture tracks the contract and fails loudly when it grows. - `tests/unit/model/deepseek/test_weight_pack_parity.py` — four properties over the real load -> pack -> stack path with the sidecar bypassed. The tests are chosen so each one fails on a specific regression the generic stacker could introduce: reproducibility (without it the rest is noise), sensitivity to layer order (a mis-ordered slab keeps every shape and dtype, so only content shows it), group provenance (perturbing a compress_ratio==4 layer must move the CSA slabs and must not touch the HCA ones, and vice versa), and the rank axis plus contiguity that `alloc_stacked_tensor` requires for the resident upload. Two things the harness had to encode rather than smooth over, both found by running it: - The compressor/indexer tensors cannot be toy-sized. The packer validates the active branch against fixed model dimensions and zero-fills the inactive branch at those same dimensions, and every slab is allocated from layer 0's template — so a toy-sized CSA tensor cannot match the template placeholder and the stack fails on a shape mismatch. The fixture imports those constants instead of copying them, and stores the two tensors the packer transposes in their pre-transpose orientation. - Each build needs its own directory. With a shared one the second checkpoint overwrote the first's shards, so the order-sensitivity test compared a checkpoint with itself and passed while proving nothing. That is exactly the failure mode a parity harness must not have, and it is why the sensitivity tests exist at all. Not asserted here: "the stacker never calls `torch.cat`". Banning it across this path would also ban `deepseek_v4_hadamard_idx`, which legitimately builds its matrix with `torch.cat` for every CSA layer; the ban only makes sense where the packer is mocked out, which is where `test_model_components.py` already keeps it. 191 tests pass across `tests/unit`.
…table Step 3 of hw-native-sys#163, first slice, and additive: nothing calls the new path in production yet. What lands is the schema, the evaluator, the rank policy, DeepSeekV4's core rule table, and the differential test that proves the two agree — so the rewire that comes next is a delegation rather than a leap. - `common/weights/spec.py` — `LayerContext` and `LayerWeightRule`. Only `transpose` and `reshape_groups` are expressible as data, because both are re-orientations of the same bytes; anything computing new values stays in a family's code rather than hiding arithmetic behind a field that looks declarative. - `common/weights/shard.py` — `Replicate`, holding both replication paths: into a slab slice (which is what keeps the stacker off `torch.cat`) or into a fresh expanded buffer. - `common/weights/packer.py` — `pack_layer`, which walks the rules in order and skips a rule whose destination is absent, so a caller can stage one group of a multi-group layout without packing the rest. - `deepseek/weight_spec.py` — the 25 core weights as data, present for every layer whatever its attention kind. The optional compressor/indexer branches, the router and the routed experts stay in code: each needs a rule kind this schema does not express yet, and inventing those kinds before their parity test exists is how a byte-identity regression gets in. Rule order is contract, not style: the slab allocator lays out whole-model tensors in the order the packed mapping is built, and the prepacked sidecar records the resulting name-to-offset map, so reordering entries would invalidate every sidecar on disk. The first test therefore asserts the rule order equals the hand-written table's order. The other four are the differential hw-native-sys#163 asks for — old and new importable side by side, compared directly rather than against a golden file: - byte-for-byte equality on the direct path; - byte-for-byte equality on the destination path, which the legacy packer needs full destinations for, so the reference pack supplies them and the comparison is taken on the subset the rules cover; - direct vs destination agreement, which is the trap hw-native-sys#163 names explicitly: the direct path casts with `.to(dtype)` while the destination path casts implicitly inside `copy_()`. They agree for the conversions in use, and this asserts it instead of assuming it; - the subset-of-destinations behaviour above. 196 tests pass across `tests/unit`.
…eights Finishes the declarative half of hw-native-sys#163 step 3. Still additive — production calls the hand-written packer — but the table is now whole, so the rewire that follows can be a single delegation instead of a half-declarative intermediate state nobody can review. Three rule kinds the core slice did not need, each added because a weight class exists that cannot be honestly expressed without it: - `OptionalWeightRule` — compressor/indexer weights, present for one attention kind and **zero-filled at a fixed shape** for the others. The inactive branch is written, not skipped: every layer must present the same kernel signature. The shape comes from the model, not the checkpoint, which is also why a synthetic checkpoint has to use production dimensions for these. - `DefaultedWeightRule` — router weights the checkpoint may omit. `required_when` names the context flag that decides whether an absent source is an error or simply "this layer does not use that router mode", so the requirement cannot quietly decay into a zero fill. - `ExpertWeightRule` + `ExpertParallel` — routed experts, sharded rather than replicated. Rank ownership is injected from the loader's own `deepseek_v4_local_expert_ids` rather than reimplemented as "contiguous block": the placement is a model decision, and a policy that assumed it would silently reshuffle a family that numbers experts differently. `SyntheticWeightRule` covers the Hadamard index, whose factory is looked up by key so the rule stays data. It sits between `csa_weights_proj` and `csa_inner_wkv` because that is where the hand-written table puts it, and the order is the sidecar's offset map. Dimensions are literals or the name of a `LayerContext` field (`resolve_shape`), which is what keeps a config-dependent shape — `gate_bias` at `n_routed_experts` — expressible as data rather than as a lambda smuggled into a table. The differential now covers the whole contract: all 49 names, parametrised over the three attention kinds, on both the direct and the destination path, plus the router case where the layer carries `gate_bias` and the `tid2eid` placeholder is the interesting half, plus a test that a required source still raises when absent. Name order is asserted too, not just contents. 204 tests pass across `tests/unit`.
Completes hw-native-sys#163 step 3: `pack_deepseek_v4_layer_weights` now builds its 49 weights by evaluating the rule table instead of by hand. Signature, output names and order, dtypes and diagnostics are unchanged — the last of those deliberately, since the wording is what users grep for, so the family's two "missing raw DeepSeekV4 ..." messages became templates the evaluator formats rather than strings buried in it. The hand-written implementation stays, renamed `_pack_deepseek_v4_layer_weights_legacy` and documented as the parity reference until hw-native-sys#163's cleanup step. That is not hesitation: the differential suite compares two independent implementations, and pointing both of its sides at the public function would have made every assertion in it tautologically true. Deleting it belongs with the step that also deletes the suite. Two tests were added because the rewire creates a gap the existing ones cannot see. The wrapper picks the context, the rank policy, the expert policy and the synthetic factories itself, so a wrong expert policy or a forgotten factory would yield plausible output that a test calling `pack_layer` directly never exercises. The new cases compare the public entry point against the legacy packer across all three attention kinds, and assert the family's error wording survives. 208 tests pass across `tests/unit`.
hw-native-sys#163 step 4. `load_stacked_layer_weights` now drives `common/weights/stacker.py` instead of three DeepSeek-shaped helpers, and the slab geometry becomes data: a `StackGroup` is a set of weights plus the *ordered* layer ids that contribute to it, and a layer's position in that list is the slice it occupies. That is what lets a group cover a subset of layers — the CSA group holds its layers contiguously in first-appearance order, which is what the fused kernels index and what the prepacked sidecars on disk were written from. Two properties of the old code kept deliberately, because both are load bearing: - **Nothing calls `torch.cat`.** Slabs are allocated once from the template layer's shapes and every later layer is packed straight into a view of its own slice. Concatenating instead would hold the whole model twice at the peak — ~346 GB for DeepSeek V4 that it does not have to pay. - **A group with no layers allocates nothing**, so a model using none of an attention kind reserves no slabs for it. The FWD group declares `members=None`, meaning "everything the other groups do not claim, in the packer's own order". A new weight therefore joins it automatically, the way `fwd_names` used to be derived, rather than needing to be listed in two places. One behavioural trap found while wiring it: moving the progress log into the pack callback silently dropped layer 0's line, because layer 0 takes the template-copy path and never reaches the packer. `stack_layers` grew an `on_layer_done` hook that fires once per layer whichever path it took. `test_weight_stacker_parity.py` compares the generic stacker against the hand-written helpers over five layouts — mixed, no groups at all, CSA-only, HCA-only, and interleaved — because placement is the half of this contract a shape check cannot see: every slab keeps its shape and dtype no matter which layer landed where, so wrong geometry produces correctly-shaped wrong weights. Each layer is given distinct values so a misplacement changes the bytes. Plus the view-aliasing property, the empty-group case, and the two rejections that keep the catch-all unambiguous. The three old helpers are unreferenced now but stay until the cleanup step, for the same reason the legacy packer does: they are the other side of this differential. 218 tests pass across `tests/unit`.
…actor The sidecars already published are the artifacts most exposed by hw-native-sys#163: each one is a whole-model stacked payload that took ~40 minutes to build, and the refactor is only safe if it neither changes what gets written nor rejects what was written before. Both directions are now covered: - the payload the current stacker writes, against the payload the hand-written helpers wrote — tensor entries, offsets, metadata values and every payload byte; - a legacy-written sidecar read back through the current reader and compared against a fresh pack of the same checkpoint, which is the direction that protects files on disk; - a currently-written sidecar actually being consumed rather than silently repacked, asserted by making the shard path fail the test if it is reached; - the fingerprint recomputed against its own definition, since it is what decides whether published files stay valid — a changed payload would invalidate every sidecar in the fleet at once, silently, by making them all look stale; - staleness still detected, via a source shard whose mtime moved. **A byte-for-byte file comparison would have been the obvious assertion and it is wrong.** safetensors serializes its metadata map in nondeterministic order: eight writes of one dict in a single process produced two different key orders, so the first version of this test failed 2 runs in 3, in the header, for a reason that has nothing to do with this refactor. Anyone "fixing" that by touching the writer would be chasing a phantom. The test now compares the header's tensor entries, the metadata as a dict, and the payload bytes — which is what a reader depends on — and the reasoning is recorded in the docstring so the trap is not rediscovered. Found by re-running the file rather than by reading it: it passed the first time. 223 tests pass across `tests/unit`, and this file five times in a row.
The two pieces hw-native-sys#163 lists that the Qwen migration cannot start without. Both are additive: DeepSeek now declares its staging policy instead of implying it, and nothing else changes behaviour. **`common/weights/pipeline.py`** owns the staging order, because peak host memory is decided there and the two families need opposite answers. DeepSeek stages serially — packing one layer allocates ~8 GB of intermediates, 256 routed experts each stacked and rank-replicated, so overlapping multiplies the peak and contends on bandwidth instead of hiding latency. Qwen wants a pool, its layers being small enough that read latency dominates. Neither is a sensible default, so `StagingPolicy` is a parameter. Two details in it are load bearing rather than cosmetic: - `workers=1` takes the no-pool path, not a pool of one, so a family that must not overlap cannot be made to by a scheduling accident. - each pooled worker pins torch to one thread for its duration and restores it after. Without that, N staging threads each fan out into torch's own pool and oversubscribe the machine — the copies then run *slower* than serially, which reads as "threading did not help" rather than as a misconfiguration. `stack_layers` now takes the policy and routes its loop through it, so the module is used rather than decorative. Overlapping is safe there by construction — each layer writes a disjoint slice of each slab — but safety is not speed, which is the caller's call to make. **`GlobalWeightRule`** covers embedding, LM head and final norm, whose three conditions are all silent when wrong: - `fallback_source` — a tied checkpoint ships no `lm_head` and the embedding stands in. Falling back is correct; inventing zeros is not. - `pad_to_multiple` — the fused LM head hard-codes a padded vocabulary, so the weight has to grow to meet it. Those rows are never selected, but they are matmul operands. - `pad_fill` — and it differs between the two: the embedding pads with zeros, the LM head pads by **replicating row 0**. Zero rows in an LM head give every padded token the same finite logit rather than an impossible one, so the mistake survives review and surfaces as sampling noise. Both behaviours are copied from the Qwen executor, which is where they exist today. Twenty tests, including the ones that would otherwise be assumed: that `workers=1` really runs on the calling thread, that a pool really overlaps (via a barrier, which a serial runner would deadlock on), that a failing layer propagates rather than leaving a half-written slab, that torch's thread count is restored, and that a layer's tensors are unreachable once its write returns — checked with weak references, since "we do not keep it" is the kind of claim that quietly stops being true. One thing removed rather than added: a `finally: raw = None` that looked like it was releasing each layer. It was theatre — the mapping is a local of one call and dies with it — so the docstring now says that instead. 243 tests pass across `tests/unit`.
hw-native-sys#163 step 5. The draft layer needed no new machinery for its 49 layer weights — `LayerContext` already carries the prefix, so `mtp.0` goes through the same table as `layers.N` — and its twelve extras, previously a hand-written dict literal, are now a rule table of their own. Ten of them are replicated reads. The two `_smooth` tensors are synthesized all-ones rows, and they carry the one trap here: the factory returns a **single row** and the rank policy replicates it to `[ranks, hidden]`. A factory returning the rank-shaped tensor directly — which is what the literal built — would come back as `[ranks, ranks, hidden]`, still all ones, so every value check would pass while the shape was wrong for the kernel. There is a test for exactly that shape. The fixture grew the draft layer, which it did not have before, so none of this was covered until now. Its names are derived the way `load_mtp_weights` derives them — `deepseek_v4_layer_weight_names` with `layers.0` rewritten — so the fixture cannot drift from the loader it is meant to exercise. Three differential tests: the twelve extras against a reproduction of the old dict literal, the synthesized pair's shape and values, and the full `load_mtp_weights` path end to end through the store, asserting the extras land last and in rule order. 246 tests pass across `tests/unit`.
hw-native-sys#163 step 6, first half: Qwen3's checkpoint contract, stacked layout and globals as data, with a differential against the executor's own staging loop. Additive — the executor still runs its loop; the metadata-only loader flip comes next. Two family differences had to become parameters of the shared code rather than assumptions inside it, and neither is cosmetic: - **Stack axis.** A DeepSeekV4 weight leads with the rank axis its upload shards on and stacks its layers on axis 1; a rank-less Qwen weight stacks on axis 0. Get it wrong and the slab is the right size holding a transposed model, which no shape check catches. `allocate_slabs` / `destinations_for` / `stack_layers` now take `stack_axis`, and use `narrow` rather than a hard-coded slice. - **No rank axis at all.** Added `NoShard` beside `Replicate`, as a distinct policy rather than `Replicate(ranks=1)`: the former yields `[*shape]` and the latter `[1, *shape]`, and a slab built from the wrong one has an extra axis nothing indexes. `allocate_slabs` also takes the allocator, because Qwen's slabs must be shared memory — its upload reads them from a forked child — while DeepSeekV4's need not be. Two schema additions, both for behaviour that is silent when wrong: - `flatten_to_row` on layer rules: a 1-D norm gamma has nothing to stack on, and reshaping it to `[1, dim]` is what makes it stackable at all. - `default_fill="ones"` on `DefaultedWeightRule`, for a Qwen3 checkpoint that ships no QK norms. That is a model variant rather than a fault, and the neutral gamma is ones — zeros would annihilate the activations the gamma scales instead of leaving them unscaled, changing the model's output without changing a single shape. - `LayerContext.dims`, a map rather than more fields, so `head_dim` is available to a rule without one family's config vocabulary accumulating on a type every family shares. **A positional-argument bug the differential caught, worth recording.** Inserting `flatten_to_row` after `default_shape` silently shifted DeepSeekV4's positional `required_when`: the string `"include_gate_bias"` landed in `flatten_to_row`, truthy, so `gate_bias` came out `[2, 1, 4]` instead of `[2, 4]` — same values, same dtype, extra axis. The field moved to the end of the dataclass and those call sites now pass by keyword, with a comment saying why. Nine Qwen tests: the eleven slabs byte-for-byte against a reproduction of `_stage_stacked_decode_weights`, absent QK norms defaulting to ones, the transposed projections and `[1, dim]` gammas, no slab carrying a rank axis, a reordered-layer sensitivity check, and the four global-weight behaviours including the tied head and the two padding fills. 255 tests pass across `tests/unit`.
…int eagerly hw-native-sys#163 step 6b: `HuggingFaceDirectoryLoader` becomes metadata-only for the per-layer weights, and the Qwen executor stages them one layer at a time through the rule table. Peak host memory drops from ~2x the model to ~1x. The eager path kept three representations alive at the peak: the whole state dict from `_load_safetensors_dir`, the cast `LayerWeights` built from it, and the stacked destinations. Now a layer's raw tensors are read, written into their slab slice, and dropped before the next layer is read; only the destinations persist. The globals stay eager, and that is a constraint rather than a choice: `Executor.lookup_embeddings` reads `model.embed_tokens` at **request** time, so the embedding cannot be released after staging. `_load_safetensors_subset` reads just those few names through `safe_open` rather than `load_file`, so reaching one tensor does not materialize the shard holding it. Removed as dead once nothing called them: `_load_safetensors_dir`, the eager `LayerWeights` loop, and `_release_layer_weights` — the last of which existed to blank out per-layer tensors after copying, which lazy reading makes unnecessary. **The staging function is now exercised end to end, which it was not before.** The earlier Qwen tests compared the rule table against a reproduction of the executor's old loop, which only proves the table matches my reading of that loop. Four new tests drive the real `_stage_stacked_decode_weights` against a Hugging Face checkpoint written to disk and compare it to that same reproduction: the eleven slabs byte-for-byte, the absent-QK-norm variant, that every slab is shared memory (the upload reads them from a forked child, so private memory would be a silent correctness bug), and that a model loaded without the metadata fails with an actionable message rather than a `KeyError`. 259 tests pass across `tests/unit`.
files the rest of it would touch. `RuntimeModel.layers` is now defaulted and empty for every loader in the tree, and `LayerWeights` is documented as deprecated with a pointer to what replaced it. Both are kept rather than deleted on purpose: hw-native-sys#168, hw-native-sys#114, hw-native-sys#152, hw-native-sys#145, hw-native-sys#144 and hw-native-sys#132 are open, and some construct a `RuntimeModel`. Removing the field would break them for no gain that cannot wait — it goes when they have landed. `--num-layers-override` in the Qwen example loses its `runtime_model.layers[:n]` slice, which had quietly become a no-op: staging reads `config.num_hidden_layers` and pulls exactly that many layers from the checkpoint, so replacing the config *is* the override now. The knob still works; there is simply nothing eager left to truncate. **The `stage_weights` hook is deliberately not added, and the reason is worth recording because the issue asks for it.** hw-native-sys#163 proposes it "between `_create_runner` and `init_kv_cache`", motivated by Qwen's `DistributedWorker` forking inside `init_kv_cache` — so staging must happen before that. Checking the actual order in `PyptoExecutor.register_model`, it already does: `_compile_model` runs first, and Qwen stages inside it, well ahead of `_create_runner` and the fork. The constraint the hook exists to satisfy is met without it. That leaves the hook as an architectural tidy-up — separating "compile kernels" from "stage weights" into named phases — and it would touch `common/runner/model_runner.py` plus both `npu_runner.py`, which is exactly where the six open PRs are. Adding surface to the most contested files in the repo to formalise a phase ordering that already holds is a poor trade this week. Worth doing after they land, with the timing rationale restated then rather than assumed.
…he new one hw-native-sys#163 step 8, which closes the issue. Net -1057 lines. **Deleted, now that the differentials that guarded them have run green in CI:** `_pack_deepseek_v4_layer_weights_legacy` and the five helpers only it used (`_replicate_weight`, `_pack_wo_a`, `_pack_deepseek_v4_optional_attention`, `_pack_deepseek_v4_router`, `_pack_deepseek_v4_routed_experts`), the three hand-written stacker helpers, and the two differential suites themselves. `weight_loader.py` goes from 1589 to 1097 lines. Keeping them longer was tempting and wrong: a differential compares two implementations, so the moment one is deleted the suite is comparing nothing. The proof is in the record — those tests ran green against the rule table on hw-native-sys#174 and on this branch — and `git log` has both sides if a byte-level question ever comes back. What remains is the property that outlives the refactor: the sidecar written by this writer round-trips through the reader and matches a fresh pack of the same checkpoint. **`LayerWeights` and `RuntimeModel.layers` are gone for real**, not just deprecated. The reason I gave for deferring this — that six open PRs construct a `RuntimeModel` — did not survive checking: none of them pass `layers=`. hw-native-sys#114's two matches are `num_layers=model.config.num_hidden_layers`, which is unrelated. Five call sites in tests and the loader dropped the argument. **Docs.** A new `docs/dev/model/weight-staging.md` covers the pipeline: the four stages, the six rule kinds, and the invariants that are not obvious from the code — rule order being sidecar contract, why nothing in the stacker calls `torch.cat`, why the stack axis differs per family, why DeepSeek stages serially and Qwen pooled, and why an LM head cannot pad with zeros. Both model docs gain a section pointing at it with their own specifics. The two traps found by running the tests rather than reading them are recorded there too, so they are not rediscovered: safetensors serializing its metadata in nondeterministic order (which makes a whole-file assertion flaky), and a synthetic checkpoint needing production dimensions for the compressor/indexer weights. 228 tests pass across `tests/unit`; `test_model_components.py` passes unmodified, which is the gate hw-native-sys#163 sets.
…n common Review on hw-native-sys#174 asked for a common base class for the Qwen and DeepSeekV4 loaders, which is the same move this PR made for the *stores* and had not made for the loaders. `SafetensorsDirectoryLoader` now holds the parts that are genuinely shared: `supports_format` was byte-identical in both, and both must find a `config.json` before any of their own checks mean anything — so `can_load` becomes that precondition plus a `_recognises` hook each family implements. `load` deliberately stays out of it. The two produce different things from a directory — one stages Hugging Face layers lazily, the other validates a quantized checkpoint contract and hands back metadata — and a shared skeleton would need a hook per step, which is more coupling than the duplication it removes. The DeepSeekV4 detection keeps its own shape rather than being folded into the base: it requires a shard index and a config that names DeepSeekV4, and treats an unreadable config as "not mine" rather than an error, so the registry moves on to the next loader. That asymmetry with the Hugging Face loader (which accepts any directory carrying safetensors) is intentional and now has a comment saying so. Also rewrote `pack_deepseek_v4_layer_weights`'s docstring, which the same review read as a no-op wrapper — fairly, because it led with the history of the hand-written packer it replaced instead of what it does. It is the family's binding to the generic evaluator: it builds the `LayerContext`, selects the rank and expert-placement policies, supplies the synthetic-weight factories, and carries the error templates its callers' diagnostics are matched against. Both production callers would otherwise repeat that wiring, and the MTP path would additionally have to thread `prefix="mtp.0"` through it. 257 tests pass; `ruff` and the repo hooks clean.
c5ddc89 to
119f2ef
Compare
|
Correction to my last paragraph: no PR from me is needed. The cap already landed on Rebased onto |
superxf
left a comment
There was a problem hiding this comment.
Requesting changes for one process-wide Torch thread-count race; see the inline comment. The local unit and lint suites otherwise pass.
|
|
||
| def _run(layer_id: int) -> int: | ||
| if policy.pin_torch_threads: | ||
| previous = torch.get_num_threads() |
There was a problem hiding this comment.
Restore the process-wide Torch thread count outside the worker pool
torch.set_num_threads() changes process-wide state, so saving and restoring it independently in each worker races. For example, worker A saves 8 and sets 1; worker B then saves 1; A restores 8; and B finally restores 1. The serving process is then left at one Torch thread after staging. I reproduced this deterministically with two event-ordered workers (torch_threads_after=1). The previous Qwen implementation avoided this by setting the count once before entering the pool and restoring it once after all workers had joined. Please keep the save/set/restore around the whole pool and add a test that forces overlapping workers to exit in this order.
There was a problem hiding this comment.
Confirmed and fixed in 1b7e292 — the race is real and the reasoning is exactly right. The pin now wraps the whole pool, under try/finally so a failing layer cannot leak it either, which is what the previous Qwen implementation did.
One correction to your test request, because I tried it your way first and it does not work. The interleaving cannot be forced from stage: each worker reads the count before its callback runs, so by the time an event inside stage can fire, both workers have already saved. I wrote that test with two event-ordered workers, ran it against the pre-fix code, and it passed — worthless as a regression test.
So it asserts the protocol instead: exactly one pin and one restore for the whole pool, whatever the worker count.
real_get = torch.get_num_threads
sets: list[int] = []
monkeypatch.setattr(torch, "set_num_threads", sets.append)
stage_layers([0, 1, 2, 3], stage=lambda _: None, policy=StagingPolicy(workers=2))
assert sets == [1, real_get()]Verified in both directions: against the pre-fix code it fails with 8 set_num_threads calls for 4 layers instead of 2; after the fix it passes. A second test covers the raising path.
If you would still rather have the ordered-worker version alongside it, I can add it as documentation of the scenario — but it should not be the regression test, since it does not fail on the bug.
…orker `torch.set_num_threads` is process-wide, so saving and restoring it inside each worker races with itself: worker A saves 4 and sets 1, worker B then saves *1*, A restores 4, and B -- finishing last -- restores 1. Staging then leaves the whole serving process single-threaded for everything that follows. Reported in review with a deterministic reproduction; the previous Qwen implementation set the count once before entering the pool, which this restores. The pin now wraps the pool, under try/finally so a failing layer cannot leak it either. On the test: the interleaving cannot be forced from `stage`, because each worker reads the count before its callback runs. My first attempt did exactly that and passed against the broken code -- worthless as a regression test. This asserts the protocol instead: exactly one pin and one restore for the whole pool, whatever the worker count. Verified to fail on the pre-fix code (8 set_num_threads calls for 4 layers instead of 2) and pass after. A second test covers the raising path.
1b7e292 to
c1b64c4
Compare
Steps 1-4 of #163, plus the sidecar parity: the family-neutral pieces, DeepSeekV4's layer contract expressed as data, and the whole-model slab geometry. Additive except for one delegation, and every step gated by a differential test rather than by inspection.
@ndleslx — you confirmed on the issue that you have no branch for this, so this is the whole refactor: the family-neutral pipeline, both families expressed as rule tables, Qwen's lazy staging, and the legacy path retired.
It was two PRs (#174 + #177) until #177's deletions made that a bad split — the first added a legacy rename and two differential suites that the second deleted, so a reviewer read work that cancelled. Combined here; #177 is closed as superseded, and its CI history is there if the intermediate state is ever of interest.
What is here
common/weights/store.pyLazySafetensorsStore— the checkpoint index and the grouped reads, lifted out ofDeepSeekV4WeightStore, which now subclasses itcommon/weights/spec.pyLayerContext, and five rule kindscommon/weights/shard.pyReplicateandExpertParallelcommon/weights/packer.pypack_layer, the generic per-layer evaluatorcommon/weights/stacker.pyStackGroupplusallocate_slabs/destinations_for/copy_packed_layer/stack_layersdeepseek/weight_spec.pypack_deepseek_v4_layer_weightsnow evaluates that table, andload_stacked_layer_weightsdrives the generic stacker instead of three DeepSeek-shaped helpers. Signatures, output names and order, dtypes and error wording are unchanged.Slab geometry becomes data: a
StackGroupis a set of weights plus the ordered layer ids contributing to it, and a layer's position in that list is the slice it occupies — which is how a group covers a subset of layers (CSA holds its layers contiguously in first-appearance order, which is what the fused kernels index and what the sidecars on disk were written from). Two old properties kept deliberately: nothing callstorch.cat, since concatenating would hold the model twice at the peak (~346 GB it need not pay), and a group with no layers allocates nothing.The parity harness comes first, on purpose
pypto-lib fuses every layer into one kernel, so the whole-model stacked layout is an output contract and "byte-identical" has to be checkable, not intended.
tests/unit/model/conftest.pyadds a(shape, dtype, sha256)fingerprint;tests/unit/model/deepseek/conftest.pywrites a synthetic checkpoint whose tensor set comes fromdeepseek_v4_layer_weight_namesrather than a copied list, so it tracks the contract instead of drifting from it.The tests are written as sensitivity rather than as green checks — each one fails on a specific regression the generic path could introduce:
compress_ratio==4layer must move the CSA slabs and must not touch the HCA ones, and vice versa;alloc_stacked_tensorneeds for the resident upload;Three things worth flagging for review
The legacy packer and the three old stacker helpers are still in the file, the packer renamed
_pack_deepseek_v4_layer_weights_legacy. Each differential compares two independent implementations; pointing both sides at the new code would make every assertion tautologically true. They go away with the suites, at the cleanup step.safetensors does not write byte-identical files, and the obvious sidecar assertion is therefore wrong. It serializes its metadata map in nondeterministic order — eight writes of one dict in a single process gave two different key orders — so a byte-for-byte file comparison failed 2 runs in 3, in the header, for a reason unrelated to this refactor. Found by re-running rather than by reading: it passed the first time. The test compares the header's tensor entries and offsets, the metadata as a dict, and the payload bytes; the reasoning is in the docstring so nobody "fixes" it by touching the writer.
One behavioural trap, found while wiring the stacker rather than by reading it: moving the progress log into the pack callback silently dropped layer 0's line, because layer 0 takes the template-copy path and never reaches the packer.
stack_layersgrew anon_layer_donehook that fires once per layer whichever path it took.The direct and destination paths cast differently — explicit
.to(dtype)on one, implicitly insidecopy_()on the other. #163 names this as a byte-identity trap. They agree for the conversions in use, and there is now a test asserting it rather than assuming it.The compressor/indexer tensors use production dimensions in the fixture. The packer validates the active branch against fixed model dimensions and zero-fills the inactive branch at the same sizes, and slabs are allocated from layer 0's template — so a toy-sized CSA tensor cannot match the placeholder. The fixture imports those constants instead of copying them.
What it buys
Peak host memory during Qwen3 startup drops from ~2× the model to ~1×. The eager path kept three representations alive at the peak — the whole state dict from
_load_safetensors_dir, the castLayerWeightsbuilt from it, and the stacked destinations. Now each layer's raw tensors are read, written into their slab slice, and dropped before the next layer is read; only the destinations persist.What could not move, and why
Executor.lookup_embeddingsreadsmodel.embed_tokensat request time, not at compile time, so the embedding cannot be released after staging — the issue flags this and it holds. The globals therefore stay eager._load_safetensors_subsetreads just those few names throughsafe_openrather thanload_file, so reaching one tensor does not materialize the whole shard around it.RuntimeModel.layersis now left empty for this loader, and the executor reads what it needs throughextra(model_dir,weight_map). Nothing else in the tree readlayers— that was checked, not assumed.Removed as dead
_load_safetensors_dir, the eagerLayerWeightsloop, and_release_layer_weights— the last existed to blank out per-layer tensors after copying, which lazy reading makes unnecessary.Step 8: the legacy path is gone
Net −1057 lines. Deleted now that the differentials guarding them have run green here and on #174:
_pack_deepseek_v4_layer_weights_legacyplus the five helpers only it used, the three hand-written stacker helpers, and the two differential suites themselves.weight_loader.pygoes from 1589 to 1097 lines.Keeping the differentials longer was tempting and wrong: they compare two implementations, so the moment one is deleted the suite compares nothing. The proof is in the record and
git loghas both sides. What remains is the property that outlives the refactor — the sidecar this writer produces round-trips through the reader and matches a fresh pack of the same checkpoint.LayerWeightsandRuntimeModel.layersare gone for real. I had deferred that on the grounds that six open PRs construct aRuntimeModel; checking rather than assuming, none of them passlayers=(#114's two matches arenum_layers=model.config.num_hidden_layers). Five call sites dropped the argument.New
docs/dev/model/weight-staging.mdcovers the pipeline and, more usefully, the invariants that are not visible in the code: rule order being sidecar contract, why nothing in the stacker callstorch.cat, why the stack axis differs per family, why DeepSeek stages serially and Qwen pooled, and why an LM head cannot pad with zeros. Both model docs point at it with their own specifics.Step 7, and one thing the issue asks for that I did not do
RuntimeModel.layersis now defaulted and empty for every loader, andLayerWeightsis documented as deprecated with a pointer to what replaced it. Both are kept rather than deleted: #168, #114, #152, #145, #144 and #132 are open and some construct aRuntimeModel, so removing the field would break them for a gain that can wait.--num-layers-overrideloses itsruntime_model.layers[:n]slice, which had quietly become a no-op — staging readsconfig.num_hidden_layersand pulls that many layers from the checkpoint, so replacing the config is the override now.The
stage_weightshook is deliberately absent. #163 proposes it between_create_runnerandinit_kv_cache, because Qwen'sDistributedWorkerforks insideinit_kv_cacheand staging must precede that. Checking the real order inPyptoExecutor.register_model:_compile_modelruns first and Qwen stages inside it, well ahead of_create_runnerand the fork — so the constraint the hook exists to satisfy already holds.That leaves it as an architectural tidy-up, separating "compile kernels" from "stage weights" into named phases, and it would touch
common/runner/model_runner.pyplus bothnpu_runner.py— exactly where those six PRs are. Adding surface to the most contested files in the repo to formalise an ordering that already holds seemed the wrong trade; happy to be told otherwise, and it is a small change to make once they land.Verification
223 tests pass across
tests/unit;test_model_components.pypasses unmodified, which is the gate #163 sets for this step.ruff,check-headersandcheck-english-onlyclean.No perf change: legacy vs rules on the same input is −1.8% to −8.0% (i.e. marginally faster, at the noise level) across the three attention kinds. That is only the dispatch overhead — the byte movement is identical by construction, since both go through the same policy.
The hardware gates (DeepSeek 8-card, Qwen3 accuracy) are not run locally, but they are meaningful again now that #173 pins the pypto revision CI builds against — that pin sits before the upstream stall (
hw-native-sys/pypto#2354), so a red gate on this PR would be about this PR. This was a draft until #173 landed precisely because that was not true.Net effect
26 files, +3337/-764.
weight_loader.pygoes from 1589 to 1097 lines;LayerWeights,RuntimeModel.layers, the hand-written packer and the three stacker helpers are gone.224 unit tests pass (four fewer than before the rebase: main's #176 deleted
test_offline.pyandtest_device_sampling.pyalong withLLMEngine).CI: the DeepSeek prefix-cache failure is not from this branch
If
unit-testscomes back red ontest_deepseek_v4_http_completion_matches_expected_text[k1-prefix-cache], that is #183, not this PR. That case issues two identical requests and asserts the second returns the same text as the first; the direction of the mismatch flips between runs, main failed the same assertion on65b69fefand went green on9dc3bf8ewithout anyone touching that path, and the base rate across recent runs is 3 in 9.It cannot be this branch: the assertion compares two responses from the same weights, so a staging bug would make both wrong and the
expected_textchecks — including the non-prefix-cache K=1 case in the same run — would fail too. They pass.