Fix chunked loss initialization in TorchFT trainer - #4575
Open
jojoinfra wants to merge 1 commit into
Open
Conversation
1. Share Trainer setup so TorchFT binds lm_head and skips the model projection. 2. Add one focused CPU regression test for non-pipeline training.
jojoinfra
requested review from
fegin,
tianyu-l,
wconstab and
wwwjn
as code owners
September 10, 2026 08:11
meta-cla
Bot
added
the
CLA Signed
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 join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
While testing DeepSeek-V4 Flash training with TorchTitan and TorchFT across 128 dies, I ran into several issues and fixed them in my downstream setup. I'd like to contribute these fixes upstream as small, focused PRs. This PR addresses one of them.
With chunked loss enabled,
FaultTolerantTrainerdoes not bind the model'slm_headtoChunkedLossWrapperor set_skip_lm_head. These steps already exist inTrainer, but are missing from TorchFT's separate initialization path.This PR extracts the existing setup into
Trainer._configure_chunked_loss()and calls it from both trainers after model construction. The helper preserves the existing handling of non-pipeline training and pipeline stages, while keeping this fix limited to chunked-loss initialization.Validation:
lm_headbinding and_skip_lm_headin the non-pipeline case: passed.