[Fix] Reasonable loss for non-distributed training #242
Conversation
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
meta-cla
bot
added
the
CLA Signed
joecummings
force-pushed
the
compare-against-hf-trainer
branch
from
1148c66 to
1846b85
Compare
joecummings
force-pushed
the
compare-against-hf-trainer
branch
from
1846b85 to
c0c0a35
Compare
| async def setup(self): | ||
| engine_config = {f.name: getattr(self, f.name) for f in fields(self)} | ||
| self.engine = ForgeEngine(ForgeJobConfig(**engine_config)) | ||
| self.engine.checkpointer.load() |
There was a problem hiding this comment.
I was actually staring at this in the trainer side .. It's unclear at a glance how the checkpointer.load is associated with loading the HF model weights
There was a problem hiding this comment.
Yeah it's not very clear without digging into the TorchTitan checkpointing code here: https://github.com/pytorch/torchtitan/blob/5b5d46856b400c8550989415bee91473aab4f921/torchtitan/components/checkpoint.py#L523
All the information is taken from the config and instantiated into the CheckpointManager. Then the load call only takes a "step", which in our case isn't needed b/c it should be a static model every time.
Jack-Khuu
reviewed
Sep 26, 2025
| async def setup(self): | ||
| engine_config = {f.name: getattr(self, f.name) for f in fields(self)} | ||
| self.engine = ForgeEngine(ForgeJobConfig(**engine_config)) | ||
| self.engine.checkpointer.load() |
allenwang28
approved these changes
Sep 26, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 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.
Context
This PR serves to fix an issue we were seeing wherein the first loss step was astronomically high (>500). The issue was tracked to a very high KL divergence value, which just measures the difference in logprobs between the reference model and the training model. For the very first step, this definitely shouldn't be the case b/c they are both the same model at that time!
Therefore, perhaps the weights on the reference model and training model were actually not the same. Further comparison against a forward pass of the Hugging Face model confirmed this hypothesis.
Fix
Load in the reference model weights correctly through the TorchTitan APIs.
Before
After
To-dos