[pull] master from deepmodeling:master - #316
Merged
Merged
Conversation
…through pt_expt assembly (#5960) Users reported that compiled DPA4 training runs ~2x slower on `pt_expt` than on `pt`. This PR is the result of chasing that: a performance bug in how the SO3 contractions were lowered, and a correctness bug where a configured `use_amp: false` was lost while `pt_expt` assembled the model. ## Changes **1. Weight broadcast across the node axis (`so3.py`, `lora.py`, `grid_net.py`)** `matmul(x[..., None, :], weight[None, ...])` makes the node count `N` the matmul BATCH, so matmul broadcasts the weight to `(N, D, F, Cin, Cout)` and autograd then reduces that whole expanded gradient (`ExpandBackward0`) back to the parameter shape. At the water example's sizes a 165 K-element weight expanded to 191 M elements (~0.8 GB) per call, and the reduce was the single costliest kernel of a training step (45.6 ms, 3x per step). The fix batches the contraction over the small `(D, F)` axes so `N` stays the GEMM ROW dimension and the weight is used in place. Micro-benchmark, fwd+bwd at the real shapes: **16.48 ms -> 1.09 ms (15x)**. Two lookalike sites in `projection.py` are deliberately NOT changed: their operands are `requires_grad=False` buffers, so no backward reduce exists. Verified rather than assumed. **2. The same lowering for the frame mixers (`_degree_batched_matmul`)** Review found the `FrameContract` / `FrameExpand` mixers still on the broadcast spelling. They now share one helper, `_degree_batched_matmul`, written identically on the dpmodel side (`dpmodel/descriptor/dpa4_nn/grid_net.py`) and the pt side (`pt/model/descriptor/sezm_nn/grid_net.py`). Because it does no reshape, an empty node axis (`N == 0`) flows through unchanged instead of hitting a reshape error — pinned by a test. **3. `use_amp` was lost during `pt_expt` model assembly (correctness)** `use_amp` is a training-runtime policy, not model state, so it stays OUT of the portable serialization record (this is the #5963 position, and the jax deserializer actively rejects records carrying `use_amp: true`). The bug was elsewhere: `pt_expt` assembled its model by converting an already-populated dpmodel instance, and that conversion round-trips the component through `deserialize(serialize())`. Anything that is deliberately not in the portable record — `use_amp` among it — was therefore dropped, and training silently ran under bfloat16 autocast even when the input configured `use_amp: false`. The fix is at the assembly boundary, not in the record: `pt_expt` now constructs the wrapped class directly (`auto_wrapped_class(...)` in `make_model.py`, `get_model.py`, and the bridging composition path), so a live constructor-supplied component keeps its runtime state. The rule is stated once, in the `auto_wrapped_class` docstring; the call sites reference it. An earlier revision of this PR instead added `use_amp` to `serialize()`. That was reverted in review — it put a runtime knob into the portable record and would have broken the jax contract. **Also removed in review: an `enable_tf32` / `DP_TF32_INFER` implementation for `pt_expt`.** It contributes nothing to the speedup measured below (the benchmark card has no TF32 silicon), and #5958 owns the `pt_expt` training-runtime alignment — including the documented position that `pt_expt` always runs at `"highest"` matmul precision. `pt_expt` therefore keeps master's warn-and-ignore behavior for `enable_tf32`. ## Benchmark DPA4 water example (`examples/water/dpa4`), one Tesla T4, torch 2.11, fp32 (`use_amp: false`), batch size 6. Steady-state seconds per training step, obtained by differencing the wall time of a 33-step and a 3-step run of the same config, which cancels every one-time cost (import, data load, statistics, `torch.compile` / make_fx lowering). All five arms were measured in one session on the same machine; run-to-run variation is about 2-3%. **Provenance: measured at `ae720432b`**, the head at which this PR was opened — i.e. BEFORE the review changes (change 2, the frame-mixer lowering, and change 3's move from `serialize()` to the assembly boundary). Change 1, which is where the entire speedup comes from, is unmodified since. The numbers have not been re-measured on the current head; a re-run is pending and I will post it rather than silently reuse these. | training mode | `pt` (reference) | `pt_expt` at master | `pt_expt` at `ae720432b` | speedup vs master | |---|---|---|---|---| | eager | 0.891 s/step | 1.555 s/step | **0.921 s/step** | **1.69x** | | compiled | 0.545 s/step | 1.611 s/step | **0.535 s/step** | **3.01x** | This reproduces the reported issue at master — `pt_expt` compiled was 3.0x slower than `pt` compiled, and even slower than its own eager path, because the broadcast-weight contraction lowers to worse code under inductor than under eager cuBLAS. After the fix `pt_expt` is at parity with `pt`: eager within 3.4%, compiled within measurement noise. ## Known limitations - **Backward numerics are covered for the frame mixers, not for the SO3 / LoRA contractions.** `test_dpa4_frame_mixers.py` compares `_degree_batched_matmul`'s weight gradient against the pt module's at rtol/atol 1e-12. The rewritten SO3 and LoRA contractions are pinned on the forward against an explicit `einsum` reference (rtol/atol 1e-12, numpy and torch namespaces); their backward is still exercised only by tracing, not compared by value. - **The `pt` / `pt_expt` TF32 policy gap remains open.** On Ampere+ cards `pt` runs training matmuls under TF32 (`enable_tf32`, default `True`) while `pt_expt` ignores the key with a warning; the two backends are not speed-comparable there. Deferred to the #5958 training-runtime series. - **The residual compiled gap vs `pt` is not stable across sessions.** An earlier session measured `pt_expt` compiled 10.8% slower than `pt` compiled; the benchmark above measured it 1.7% faster. Both are within a couple of run-to-run standard deviations, so I treat compiled as at parity and the earlier gap as unconfirmed. - The history contains churn at the GridBranch router (`7518a417c` -> `01c58e665` -> `75459610a` -> `504bb2430` -> `157444204`): a matmul spelling introduced, reverted, reintroduced, and finally restored to master's line. The site is byte-identical to master in the final diff. The degenerate GEMM that profiling found there existed only on this branch, so it is not a fix — I have left the commits rather than rewriting pushed history, and would squash them on request. - Unrelated but found while benchmarking: **torch >= 2.11 ships no Volta (CC 7.0) kernels**, and compiled training requires >= 2.11 via `check_compile_torch_version`. Compiled DPA4 training is therefore impossible on V100 with official wheels; T4 (CC 7.5) is the oldest card that works. ## Tests - `source/tests/common/dpmodel/test_dpa4_frame_mixers.py` — `_degree_batched_matmul` vs the pt module: forward parity, the `N == 0` contract, and weight-gradient parity. - `source/tests/common/dpmodel/test_dpa4_lora.py` — new `test_lora_so3_call_matches_einsum_contract`: `LoRASO3.call` against the explicit `einsum("ndfi,difo->ndfo")` reference with a nonzero adapter, on both the numpy and torch namespaces, for `n_focus` 1 and 2. - `source/tests/pt_expt/model/test_get_model_dpa4.py` — `use_amp` survives model assembly (both branches), for the plain and the bridged/composed construction paths. - `source/tests/common/dpmodel/test_descrpt_dpa4.py` — `use_amp` is absent from the portable serialization record and defaults on deserialize. - Existing `test_grid_branch[1]`/`[2]` cover the changed SO3 contraction against the pt implementation at rtol 1e-12. - Run locally: 434 passed / 10 skipped across the dpa4 dpmodel, pt_expt and cross-backend parity suites, plus the pt_expt model suite. CUDA-gated precision-context cases were run on a T4 (29/29). ## Test status caveat — resolved An earlier revision of this description flagged two locally failing `pt_expt` AOTI-freeze tests (`test_zbl_bridging.py::test_native_spin_with_bridging_graph_freeze_and_deep_eval`, `test_dpa4_zbl_parallel.py::TestBridgedSpinGraphSelfComm::test_freeze_embeds_with_comm_artifact`) as unadjudicated. They are now adjudicated as **pre-existing and environmental, not caused by this branch**: a clean `upstream/master` worktree on the same machine fails both with the identical `InductorError: assert isinstance(index, CppCSEVariable) and index.is_vec` (torch 2.11 CPU-SIMD codegen bug on an `atomic_add` scatter buffer), and both tests pass on this branch with the known workaround `torch._inductor.config.cpp.simdlen = 1` (2 passed). The same bug is already documented in `source/tests/infer/gen_dpa4.py` / `gen_dpa2.py`. --------- Co-authored-by: Han Wang <wang_han@iapcm.ac.cn> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
See Commits and Changes for more details.
Created by
pull[bot] (v2.0.0-alpha.4)
Can you help keep this open source service alive? 💖 Please sponsor : )