DeepSpeed
60dd200a - fix: AutoTP _replace_module uses full hierarchical path for pattern matching

Commit
85 days ago
fix: AutoTP _replace_module uses full hierarchical path for pattern matching The partition_config path in _replace_module built full_name from only the immediate parent name (prev_name), not the full hierarchical path. This meant patterns like '.*\.self_attn\.q_proj\.weight$' never matched because the name passed was just 'self_attn.q_proj.weight' (missing 'model.layers.N.' prefix). Fix: use class_name (which accumulates the full path) instead of prev_name. Verified: 14B Qwen2.5 with AutoTP=2 now correctly shards from 14.77B to 8.16B params per rank, GPU mem drops from 29.7GB to 16.5GB. HybridEngineRollout decode at 26.4 tokens/s (37.9ms/step) with optimizer, no OOM on 2x RTX PRO 6000. Signed-off-by: Guokai Ma <guokai.ma@intel.com>
Author
Parents
Loading