Add n_fused as an input from native_module (#5894)

2025-06-19 12:12:46 +00:00 · 2024-07-23 11:15:39 -04:00 · 2024-07-23 11:15:39 -04:00 · a521ffc9f8
commit a521ffc9f8
parent 4ec17a7cdf
1 changed files with 1 additions and 0 deletions
--- a/colossalai/shardformer/layer/qkv_fused_linear.py
+++ b/colossalai/shardformer/layer/qkv_fused_linear.py
@ -695,6 +695,7 @@ class FusedLinear1D_Col(ParallelModule):
            process_group=process_group,
            weight=module.weight,
            bias_=module.bias,
+            n_fused=n_fused,
            *args,
            **kwargs,
        )