[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!

此拉取请求引入了一项新功能，旨在显著提升框架中损失缩放（loss_scale）机制的灵活性和可定制性。通过允许将 inputs.extra_kwargs 直接传递给 self.loss_scale 方法，用户现在可以根据数据集样本中的特定属性或任务需求，动态地调整损失缩放策略，从而实现更精细的训练控制，同时确保与现有代码的完全兼容性。

Highlights

• 功能增强: 在 swift/llm/template/base.py 中，self.loss_scale 方法现在能够接收并传递 inputs.extra_kwargs，从而允许用户通过数据集样本自定义损失缩放行为。
• 灵活性提升: 用户可以通过在数据集样本中定义 extra_kwargs 来实现更精细化的训练控制，例如基于样本难度动态调整损失权重或应用不同的缩放策略。
• 向后兼容: 此修改是非破坏性的，完全向后兼容现有代码，不会影响其正常运行。
• 代码优化: 修复了 _data_collator 方法中一个 assert 语句的字符串格式问题。

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旨在通过将inputs.extra_kwargs传递给loss_scale方法，来增强loss_scale行为的灵活性。这是一个很好的功能增强。代码修改本身是正确的，将额外的参数成功传递给了loss_scale的__call__方法。

然而，为了让这个功能完全生效，还需要在swift/plugin/loss_scale/loss_scale.py中进行相应的修改。目前的实现中，LossScale.__call__方法虽然接收了**kwargs，但并未将其传递给get_loss_scale方法。我已经在具体的代码行提出了审查意见，建议进行修改以确保extra_kwargs能够被底层的loss_scale逻辑使用。

此外，我还注意到loss_scale.py中的一些LossScale子类在覆盖get_loss_scale方法时，方法签名不一致或没有正确传递**kwargs，这可能也需要一并修复，以保证框架的健壮性和一致性。

[gemini-code-assist]

这个修改是正确的，它将extra_kwargs传递给了self.loss_scale，也就是LossScale.__call__方法。

然而，为了使这个功能完全生效，还需要修改swift/plugin/loss_scale/loss_scale.py中的LossScale.__call__方法，以确保它能将接收到的**kwargs向下传递给get_loss_scale方法。

在swift/plugin/loss_scale/loss_scale.py的LossScale.__call__实现中，调用get_loss_scale时没有传递这些额外的参数。如果不进行此项修改，您在base.py中传递的extra_kwargs将不会被实际的loss_scale计算逻辑所使用。

考虑到您在PR描述中提到已经成功测试了此功能，请确认是否遗漏了对loss_scale.py文件的提交，或者是否有其他机制来处理这些参数。