From 01ab76eb1f403cac1db73975ee0d5ba33f19d454 Mon Sep 17 00:00:00 2001 From: Insu Jang Date: Mon, 8 Jul 2024 12:34:51 +0000 Subject: [PATCH] Add n_fused as an input from native_module --- colossalai/shardformer/layer/qkv_fused_linear.py | 1 + 1 file changed, 1 insertion(+) diff --git a/colossalai/shardformer/layer/qkv_fused_linear.py b/colossalai/shardformer/layer/qkv_fused_linear.py index 0f6595a7c4d6..000934ad91a2 100644 --- a/colossalai/shardformer/layer/qkv_fused_linear.py +++ b/colossalai/shardformer/layer/qkv_fused_linear.py @@ -695,6 +695,7 @@ def from_native_module( process_group=process_group, weight=module.weight, bias_=module.bias, + n_fused=n_fused, *args, **kwargs, )