[GPT-OSS] Add sharding configs to support Qwix quantization

Signed-off-by: Jordan Dotzel <amishacorns@users.noreply.github.com>

Author: amishacorns

tpu_inference/models/jax/gpt_oss.py

@@ -82,7 +82,7 @@ def __init__(self,
             hidden_size=hidden_size,
             dtype=dtype,
             rngs=self.rng,
-            vd_sharding=(('data', 'model'), None),
+            vd_sharding=P(('data', 'model'), None),
             random_init=self.random_init,
         )
 
@@ -105,9 +105,9 @@ def __init__(self,
                 query_tnh=P(None, 'model', None),
                 keyvalue_skh=P(None, 'model', None),
                 attn_o_tnh=P(None, 'model', None),
-                dnh_sharding=(None, 'model', None),
-                dkh_sharding=(None, 'model', None),
-                nhd_sharding=('model', None, None),
+                dnh_sharding=P(None, 'model', None),
+                dkh_sharding=P(None, 'model', None),
+                nhd_sharding=P('model', None, None),
                 mesh=self.mesh,
             )
 
@@ -120,9 +120,9 @@ def __init__(self,
                 dtype=dtype,
                 router_act='softmax',
                 random_init=self.random_init,
-                activation_ffw_td=('data', None),
-                ed_sharding=('model', None),
-                e_sharding=('model', ),
+                activation_ffw_td=P('data', None),
+                ed_sharding=P('model', None),
+                e_sharding=P('model'),
             )
 
             moe_mlp = GptOssMoE(
@@ -135,10 +135,10 @@ def __init__(self,
                 router=router,
                 swiglu_limit=swiglu_limit,
                 # Sharding configuration
-                activation_ffw_td=('data', None),
-                edf_sharding=('model', None, None),
-                efd_sharding=('model', None, None),
-                ed_sharding=('model', None),
+                activation_ffw_td=P('data', None),
+                edf_sharding=P('model', None, None),
+                efd_sharding=P('model', None, None),
+                ed_sharding=P('model', None),
             )
 
             block = TransformerBlock(
@@ -148,13 +148,15 @@ def __init__(self,
                     epsilon=rms_norm_eps,
                     dtype=dtype,
                     rngs=self.rng,
+                    activation_ffw_td=P('data', None),
                 ),
                 pre_mlp_norm=RMSNorm(
                     dims=hidden_size,
                     random_init=self.random_init,
                     epsilon=rms_norm_eps,
                     dtype=dtype,
                     rngs=self.rng,
+                    activation_ffw_td=P('data', None),
                 ),
                 attn=attn,
                 custom_module=moe_mlp,
@@ -167,15 +169,16 @@ def __init__(self,
             random_init=self.random_init,
             epsilon=rms_norm_eps,
             dtype=dtype,
+            activation_ffw_td=P('data', None),
         )
 
         self.lm_head = LMhead(
             vocab_size=vocab_size,
             hidden_size=hidden_size,
             dtype=dtype,
             rngs=self.rng,
-            vd_sharding=(('data', 'model'), None),
-            dv_sharding=(None, ('data', 'model')),
+            vd_sharding=P(('data', 'model'), None),
+            dv_sharding=P(None, ('data', 'model')),
             random_init=self.random_init,
         )