Code update

wang2yn84 · The tunix Authors · commit e165692fe901 · 2025-12-23T09:12:00.000-08:00
PiperOrigin-RevId: 848200226
diff --git a/tunix/models/gemma/model.py b/tunix/models/gemma/model.py
@@ -79,7 +79,7 @@ def get_default_sharding(is_sampling: bool = False):
         ffw_weight_fd=('tp', fsdp),
         rms_norm_weight=('tp',),
         act_btd=('fsdp', None, None if is_sampling else 'tp'),
-        act_btf=('fsdp', None, None),
+        act_btf=('fsdp', None, 'tp'),
         act_btnh=('fsdp', None, 'tp', None),
         score_weight_d1=(fsdp, None),
     )
diff --git a/tunix/models/gemma3/model.py b/tunix/models/gemma3/model.py
@@ -72,7 +72,7 @@ def get_default_sharding(is_sampling: bool = False):
         ffw_weight_fd=('tp', fsdp),
         rms_norm_weight=('tp',),
         act_btd=('fsdp', None, None if is_sampling else 'tp'),
-        act_btf=('fsdp', None, None),
+        act_btf=('fsdp', None, 'tp'),
         act_btnh=('fsdp', None, 'tp', None),
     )
 
diff --git a/tunix/models/llama3/model.py b/tunix/models/llama3/model.py
@@ -73,7 +73,7 @@ def get_default_sharding(is_sampling: bool = False):
         ffw_weight_fd=('tp', fsdp),
         rms_norm_weight=('tp',),
         act_btd=('fsdp', None, None if is_sampling else 'tp'),
-        act_btf=('fsdp', None, None),
+        act_btf=('fsdp', None, 'tp'),
         act_btnh=('fsdp', None, 'tp', None),
     )
 
diff --git a/tunix/models/qwen2/model.py b/tunix/models/qwen2/model.py
@@ -75,7 +75,7 @@ def get_default_sharding(is_sampling: bool = False):
         ffw_weight_fd=('tp', fsdp),
         rms_norm_weight=('tp',),
         act_btd=('fsdp', None, None if is_sampling else 'tp'),
-        act_btf=('fsdp', None, None),
+        act_btf=('fsdp', None, 'tp'),
         act_btnh=('fsdp', None, 'tp', None),
         exp_weight_cdf=('fsdp', None, 'tp'),
         exp_weight_cfd=('fsdp', 'tp', None),
diff --git a/tunix/models/qwen3/model.py b/tunix/models/qwen3/model.py
@@ -75,7 +75,7 @@ def get_default_sharding(is_sampling: bool = False):
         ffw_weight_fd=('tp', fsdp),
         rms_norm_weight=('tp',),
         act_btd=('fsdp', None, None if is_sampling else 'tp'),
-        act_btf=('fsdp', None, None),
+        act_btf=('fsdp', None, 'tp'),
         act_btnh=('fsdp', None, 'tp', None),
         exp_weight_cdf=('fsdp', None, 'tp'),
         exp_weight_cfd=('fsdp', 'tp', None),

Original file line number	Diff line number	Diff line change
`@@ -79,7 +79,7 @@ def get_default_sharding(is_sampling: bool = False):`
`79`	`79`	`ffw_weight_fd=('tp', fsdp),`
`80`	`80`	`rms_norm_weight=('tp',),`
`81`	`81`	`act_btd=('fsdp', None, None if is_sampling else 'tp'),`
`82`		`- act_btf=('fsdp', None, None),`
	`82`	`+ act_btf=('fsdp', None, 'tp'),`
`83`	`83`	`act_btnh=('fsdp', None, 'tp', None),`
`84`	`84`	`score_weight_d1=(fsdp, None),`
`85`	`85`	`)`
Original file line number	Diff line number	Diff line change
`@@ -72,7 +72,7 @@ def get_default_sharding(is_sampling: bool = False):`
`72`	`72`	`ffw_weight_fd=('tp', fsdp),`
`73`	`73`	`rms_norm_weight=('tp',),`
`74`	`74`	`act_btd=('fsdp', None, None if is_sampling else 'tp'),`
`75`		`- act_btf=('fsdp', None, None),`
	`75`	`+ act_btf=('fsdp', None, 'tp'),`
`76`	`76`	`act_btnh=('fsdp', None, 'tp', None),`
`77`	`77`	`)`
`78`	`78`
Original file line number	Diff line number	Diff line change
`@@ -73,7 +73,7 @@ def get_default_sharding(is_sampling: bool = False):`
`73`	`73`	`ffw_weight_fd=('tp', fsdp),`
`74`	`74`	`rms_norm_weight=('tp',),`
`75`	`75`	`act_btd=('fsdp', None, None if is_sampling else 'tp'),`
`76`		`- act_btf=('fsdp', None, None),`
	`76`	`+ act_btf=('fsdp', None, 'tp'),`
`77`	`77`	`act_btnh=('fsdp', None, 'tp', None),`
`78`	`78`	`)`
`79`	`79`