Support recomputing in HybridModel - #4496
Conversation
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
There was a problem hiding this comment.
Logic looks correct — the _checkpointed_forward method properly mirrors TransformerBlock._checkpointed_forward adapted for the hybrid layer dispatch, and the test coverage (including numerical equivalence) is solid.
One minor cleanup: the last line of hybrid_block.py has trailing null bytes and a missing EOF newline (editor artifact). See inline comment.
|
/ok to test a4071bf |
a4071bf to
554e044
Compare
|
/ok to test 74cf66e |
74cf66e to
d9b3287
Compare
|
/ok to test d9b3287 |
|
/claude review |
d9b3287 to
ce3bbbc
Compare
|
/ok to test ce3bbbc |
xuantengh
changed the title
|
/ok to test b789f96 |
|
@jaredcasper kindly reminder for approval |
svcnvidia-nemo-ci
added
Approved
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25829191438 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/25836746510 |
Signed-off-by: yhgalaxy <yhgalaxy@outlook.com>
Signed-off-by: Jon Barker <jbarker@aws-cmh-slurm-1-vscode-02.cm.cluster>
Signed-off-by: Dmytro Pykhtar <dpykhtar@nvidia.com>
7 participants
This PR refactors the recompute procedure in
transformer_block.pyto share the common logic tohybrid_block.py, enabling recompute inHybridModel, including bothfullandselectivegranularity.For
HybridModel, recompute enabled runnings can produce bitwise equal loss curve compared to non-recompute running: wandb.