DeepSpeed
2d34fd3e - fix: enable CUDA graph capture with DeepSpeed hybrid engine

Commit
88 days ago
fix: enable CUDA graph capture with DeepSpeed hybrid engine - Remove DS forward hooks (tput_timer.synchronize()) during graph capture - Fix trainer: sync_weights_from_student -> sync_weights, add eval/train toggle - Wire use_graph_capture through OPSD config and rollout wrapper - Add smoke_hybrid_gc.json + smoke_ds_zero0.json configs for testing The root cause was DeepSpeedEngine._module_forward_pre_hook calling synchronize() via tput_timer.start(), which is illegal during CUDA graph capture. We temporarily remove hooks for the capture+decode section and restore them in the finally block.
Author
Parents
Loading