fix elastic code in transformers>=4.57.6#9614
Conversation
There was a problem hiding this comment.
Code Review
This pull request refactors DeepSpeed Elastic configuration setup by moving the initialization logic from a callback to argument initialization, and adds support for fixing optimizer step devices when loading universal checkpoints. The review feedback highlights several robustness improvements, including preventing a device mismatch RuntimeError in _fix_optimizer_step_device by excluding the step tensor when determining the target device, avoiding a potential TypeError if callbacks is uninitialized, guarding against a ZeroDivisionError when calculating gradient accumulation steps, and ensuring checkpoint is a dictionary before setting load_universal to prevent type errors.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
1664765 to
a80b55f
Compare
PR type
PR information
1 修复transformers>=4.57.6 时 _sorted_checkpoints 不存在问题;
2 开启deepspeed 弹性开关时,改为在_init_deepspeed 中通过prepare_deepspeed_elastic_config 方法用deepspeed 中通过compute_elastic_config 计算出来的final_batch_size, micro_batch_size 覆盖实际配置的,因为弹性扩缩容需要保证总体batch_size 不变才不影响训练效果
Experiment results
Paste your experiment result here(if needed).