[Feature] Update loss_scale method call to pass through inputs.extra_kwargs

[CJack812]

PR type

• Bug Fix
• New Feature
• Document Updates
• More Models or Datasets Support

PR information

问题描述

在 swift/llm/template/base.py 中调用 self.loss_scale 时，没有将额外的参数传递给底层的 get_loss_scale 函数，导致用户无法通过数据集样本自定义参数来灵活控制 loss_scale行为。

修改内容

文件： swift/llm/template/base.py

修改前：

res_context_list, loss_scale_list = self.loss_scale(res_context_list, res_context_types, inputs.messages)

修改后：

res_context_list, loss_scale_list = self.loss_scale(res_context_list, res_context_types, inputs.messages, **inputs.extra_kwargs)

get_loss_scale 函数签名原本就支持额外的 **kwargs 参数：

def get_loss_scale(self, context: str, context_type: ContextType, is_last_round: bool, **kwargs) -> Tuple[List[str], List[float]]:

但在template base.py的调用处没有利用这一特性，导致额外的参数无法传递

用户现在可以在数据集的样本中定义 extra_kwargs，通过这些参数来自定义loss_scale
无破坏性修改：完全向后兼容，不影响现有代码的正常运行
支持高级用法：例如：

• 基于样本难度动态调整损失权重
• 根据任务类型应用不同的缩放策略

Experiment results

已在qwen3-omni模型的megatron sft训练和qwen2.5-omni模型的deepspeed sft 训练中试验了该功能。具体为：在样本中增加extra_kwargs ，自定义loss_scale根据这些extra_kwargs 调整 loss_scale , 模型性能提升。

该修改为框架基础设施优化，不直接影响模型性能指标，但显著提升了框架的灵活性和可定制性，为用户实现更精细化的训练控制提供了可能。

[gemini-code-assist]

Summary of Changes

Hello @CJack812, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

此拉取请求引入了一项新功能，允许用户在训练过程中更精细地控制损失缩放。通过将数据集样本中定义的额外参数传递给损失缩放方法，框架的灵活性和可定制性得到了显著提升，从而支持更高级的训练策略。除了这项主要功能外，还对多个文件中的f-string语法进行了统一和修正，将双引号内的键引用改为单引号，以提高代码一致性和避免潜在的语法问题。

Highlights

• 动态损失缩放参数传递: 在 swift/llm/template/base.py 中，self.loss_scale 方法现在会传递 inputs.extra_kwargs。这允许用户通过数据集样本中的自定义参数来灵活控制损失缩放行为，例如根据样本难度或任务类型调整损失权重。
• 代码风格统一: 对多个文件中的f-string语法进行了统一和修正，将双引号内的字典键引用改为单引号，以提高代码一致性并避免潜在的语法问题。

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

本次PR的核心目标是向loss_scale方法传递extra_kwargs，这是一个非常有价值的改进，增强了训练的灵活性，实现得很好。

然而，除了核心功能外，本次PR还包含了大量跨多个文件的代码格式化改动。虽然这有助于代码风格的统一，但也为代码审查带来了许多噪音。建议未来将功能性变更和代码格式化分开提交。

最重要的问题是，PR中新增了一个名为stepfun.ipynb的Jupyter notebook文件。此文件似乎与本次PR的主要功能无关，并且其输出单元格中包含ModuleNotFoundError错误。通常不建议将带有输出（尤其是错误信息）的notebook文件提交到代码库中。请考虑从本次合并请求中移除此文件。

[gemini-code-assist]

这个新添加的 Jupyter notebook (stepfun.ipynb) 似乎与本次 PR 的核心功能无关。它包含了带有 ModuleNotFoundError 错误的输出单元格，这表明它可能是一个实验性文件或存在未声明的依赖。通常不建议将带有输出（尤其是错误）的 notebook 文件提交到代码库中。请考虑从本次合并请求中移除此文件，以保持代码库的整洁，并在未来的独立 PR 中提交（如果需要的话）。

[gemini-code-assist]

核心修改正确，通过 **inputs.extra_kwargs 将额外参数传递给 loss_scale 函数，增强了框架的灵活性。但除了这个核心改动，本次PR在大量文件中引入了许多仅涉及引号替换的格式化修改，这增加了代码审查的噪音。建议未来将功能性修改和纯格式化修改分在不同的PR中提交。