Skip to content

[bugfix] fix mllm megatron gkd sft_alpha #8394

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
hjh0119 merged 1 commit into from
Mar 23, 2026
Merged

[bugfix] fix mllm megatron gkd sft_alpha #8394

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions swift/megatron/trainers/gkd_trainer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -664,8 +664,10 @@ def loss_func(self,
if self.sft_alpha > 0 and data_source != DataSource.STUDENT:
args = self.args
logits_sbv = student_logits.transpose(0, 1).contiguous()
per_token_loss = self.unwrapped_models[0].compute_language_model_loss(labels, logits_sbv)

model = self.unwrapped_models[0]
if hasattr(model, 'language_model'):
model = model.language_model
per_token_loss = model.compute_language_model_loss(labels, logits_sbv)
Comment on lines +667 to +670
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot Mar 21, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

For improved readability and conciseness, you can use getattr with a default value. This avoids reassigning the model variable and makes the logic for obtaining the correct model more direct and idiomatic.

            language_model = getattr(self.unwrapped_models[0], 'language_model', self.unwrapped_models[0])
            per_token_loss = language_model.compute_language_model_loss(labels, logits_sbv)

loss_mask = labels != -100
sft_loss_sum = (per_token_loss * loss_mask).sum()
sft_loss_count = loss_mask.sum().float()
Expand Down
Loading