refactor(loss): migrate DPOLossConfig, DistillationLossConfig, DraftC…#2520
Open
NolenLiang wants to merge 4 commits into
Open
refactor(loss): migrate DPOLossConfig, DistillationLossConfig, DraftC…#2520NolenLiang wants to merge 4 commits into
NolenLiang wants to merge 4 commits into
Conversation
…rossEntropyLossConfig to BaseModel Convert 3 TypedDict loss config classes to pydantic BaseModel with extra="allow". Update dict-style access (cfg["key"]) to attribute access (cfg.key) in their __init__ methods. - DPOLossConfig: 5 required fields, used by DPOLossFn - DistillationLossConfig: 3 required fields, used by DistillationLossFn - DraftCrossEntropyLossConfig: 1 optional field (arbitrary_types_allowed for ProcessGroup) Signed-off-by: nliang <nliang@nvidia.com>
NolenLiang
added
the
CI:Lfast
Contributor
Author
|
/ok to test c340a3b |
Callers (dpo.py, tests) pass DPOConfig dict or plain dict to DPOLossFn, not DPOLossConfig BaseModel. Add isinstance check to auto-convert dict to BaseModel, maintaining backward compatibility. Same fix for DistillationLossFn. Signed-off-by: nliang <nliang@nvidia.com>
Contributor
Author
|
/ok to test b5427ff |
Contributor
Author
|
/ok to test b5427ff |
NolenLiang
added
CI:L1
Add defaults to DPOLossConfig and DistillationLossConfig fields matching the reference configs (dpo.yaml, distillation_math.yaml). Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: nliang <nliang@nvidia.com>
Contributor
Author
|
/ok to test 857f01e |
Contributor
Author
|
/ok to test 5178546 |
yuki-97
reviewed
May 19, 2026
Comment on lines
+865
to
+867
| def __init__(self, cfg: DPOLossConfig | dict, use_linear_ce_fusion: bool = False): | ||
| if isinstance(cfg, dict): | ||
| cfg = DPOLossConfig(**cfg) |
Contributor
There was a problem hiding this comment.
let's not use this tricky way and fix the places that fail because of this.
Suggested change
| def __init__(self, cfg: DPOLossConfig | dict, use_linear_ce_fusion: bool = False): | |
| if isinstance(cfg, dict): | |
| cfg = DPOLossConfig(**cfg) | |
| def __init__(self, cfg: DPOLossConfig, use_linear_ce_fusion: bool = False): |
Contributor
Author
There was a problem hiding this comment.
On this branch dpo.py is not modified — DPOLossFn(master_config.dpo) at dpo.py:270 still passes a plain dict (DPOConfig is still a TypedDict here). Removing the guard now would break the L1 functional test. Will remove it once DPO PR #2524 merges.
Comment on lines
+972
to
+974
| def __init__(self, cfg: DistillationLossConfig | dict): | ||
| if isinstance(cfg, dict): | ||
| cfg = DistillationLossConfig(**cfg) |
Contributor
There was a problem hiding this comment.
Contributor
Author
There was a problem hiding this comment.
Done. Removed the guard and updated all callers in test_loss_functions.py and test_distillation.py to pass DistillationLossConfig(...) directly.
Comment on lines
+28
to
+30
| class DraftCrossEntropyLossConfig(BaseModel, extra="allow"): | ||
| model_config = {"arbitrary_types_allowed": True} | ||
| vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None |
Contributor
There was a problem hiding this comment.
I think let's not change this for now, actually DraftCrossEntropyLossConfig is not used and not set from the config.
Suggested change
| class DraftCrossEntropyLossConfig(BaseModel, extra="allow"): | |
| model_config = {"arbitrary_types_allowed": True} | |
| vocab_parallel_group: Optional[torch.distributed.ProcessGroup] = None | |
| class DraftCrossEntropyLossConfig(TypedDict): | |
| vocab_parallel_group: Optional[torch.distributed.ProcessGroup] |
Contributor
Author
There was a problem hiding this comment.
Done. Reverted to TypedDict since it is unused and not loaded from configuration files.
… DistillationLossFn guard 1. Revert DraftCrossEntropyLossConfig to TypedDict (unused, not loaded from config) 2. Remove isinstance(cfg, dict) guard from DistillationLossFn.__init__ and update all callers to pass DistillationLossConfig directly 3. Keep DPOLossFn guard for now (dpo.py still passes dict on this branch; will remove after DPO PR #2524 merges) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: nliang <nliang@nvidia.com>
NolenLiang
force-pushed
the
nliang/typeddict-to-basemodel-loss
branch
from
5178546 to
fd0af2d
Compare
Contributor
Author
|
/ok to test fd0af2d |
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.
…rossEntropyLossConfig to BaseModel
Convert 3 TypedDict loss config classes to pydantic BaseModel with extra="allow". Update dict-style access (cfg["key"]) to attribute access (cfg.key) in their init methods.
What does this PR do ?
Add a one line overview of what this PR aims to accomplish.
Issues
List issues that this PR closes (syntax):
Usage
# Add a code snippet demonstrating how to use thisBefore your PR is "Ready for review"
Pre checks:
Additional Information