wip

Author: wenxindongwork

tpu_inference/layers/vllm/fused_moe.py

@@ -103,6 +103,7 @@ def tensor_sharded_gmm_merged_column_parallel(
     rhs: jax.Array,
     rhs_bias: jax.Array | None,
     group_sizes: jax.Array,
+    group_sizes_global: jax.Array,
     transpose_rhs: bool,
     mesh: Mesh,
     intermediate_size: int,
@@ -128,9 +129,9 @@ def tensor_sharded_gmm_merged_column_parallel(
         check_rep=False,
     )(lhs, rhs, group_sizes)
 
+    
     if rhs_bias is not None:
-        rhs_bis = jnp.repeat(rhs_bias, group_sizes, 0, total_repeat_length=m)
-        # Maybe need to add sharding constraint here
+        rhs_bis = jnp.repeat(rhs_bias, group_sizes_global, 0, total_repeat_length=m)
         gmm_result = (gmm_result + rhs_bis).astype(gmm_result.dtype)
 
     n_shards = mesh.shape['model'] * mesh.shape.get('attn_dp', 1)
@@ -145,6 +146,7 @@ def tensor_sharded_gmm_row_parallel(
     rhs: jax.Array,
     rhs_bias: jax.Array | None,
     group_sizes: jax.Array,
+    group_sizes_global: jax.Array,
     transpose_rhs: bool,
     mesh: Mesh,
 ) -> jax.Array:
@@ -173,11 +175,9 @@ def _gmm_all_reduce(lhs, rhs, group_sizes):
     out_specs=(P(ShardingAxisName.MLP_DATA)),
     check_rep=False,
     )(lhs, rhs, group_sizes)
-
+    jax.debug.print("gmm_result before bias {} {}", gmm_result.sum(), gmm_result.ravel()[:10])
     if rhs_bias is not None:
-        
-        rhs_bias = jnp.repeat(rhs_bias, group_sizes, 0, total_repeat_length=m)
-        # wenxindong: Maybe need to add sharding constraint here
+        rhs_bias = jnp.repeat(rhs_bias, group_sizes_global, 0, total_repeat_length=m)
         gmm_result = (gmm_result + rhs_bias).astype(gmm_result.dtype)
 
     return gmm_result
@@ -365,6 +365,7 @@ def fused_moe_func(
     gating_output = jax.lax.with_sharding_constraint(
             gating_output, NamedSharding(mesh, P(ShardingAxisName.ATTN_DATA, None)))
 
+    jax.debug.print("hidden_state before MoE {} {}", hidden_states.sum(), hidden_states.ravel()[:10])
     hidden_states = hidden_states.reshape(num_tokens, hidden_size)
     gating_output = gating_output.reshape(num_tokens, global_num_experts)
 
@@ -381,19 +382,25 @@ def _process_tokens_locally(hidden_states_local, topk_indices_local):
         topk_argsort_revert_indices = jnp.argsort(topk_argsort_indices)
         token_indices = jnp.arange(num_tokens_local, dtype=jnp.int32).repeat(topk)
         token_indices_sorted = token_indices[topk_argsort_indices]
-        group_sizes = jnp.bincount(topk_indices_flat, length=global_num_experts)
+        group_sizes_local = jnp.bincount(topk_indices_flat, length=global_num_experts)
+        
+        # Reduce group_sizes once across data parallel shards to get global counts
+        # This is needed for bias addition and should be done only once for efficiency
+        group_sizes_global = jax.lax.psum(group_sizes_local, axis_name=ShardingAxisName.MLP_DATA)
 
         x = hidden_states_local[token_indices_sorted]
-        return x, group_sizes, topk_argsort_revert_indices
+        return x, group_sizes_local, group_sizes_global, topk_argsort_revert_indices
 
-    x, group_sizes, topk_argsort_revert_indices = shard_map(
+    x, group_sizes, group_sizes_global, topk_argsort_revert_indices = shard_map(
         _process_tokens_locally,
         mesh=mesh,
         in_specs=(P(ShardingAxisName.ATTN_DATA, None), P(ShardingAxisName.ATTN_DATA, None)),
-        out_specs=(P(ShardingAxisName.ATTN_DATA, None), P(), P(ShardingAxisName.ATTN_DATA)),
+        out_specs=(P(ShardingAxisName.ATTN_DATA, None), P(ShardingAxisName.ATTN_DATA), P(), P(ShardingAxisName.ATTN_DATA)),
         check_rep=False,
     )(hidden_states, topk_indices)
-
+    
+    jax.debug.print("hidden_state before gmm {} {}", x.sum(), x.ravel()[:10])
+    jax.debug.print("group_sizes {} {}", group_sizes.sum(), group_sizes)
     if use_ep:
         x = expert_sharded_gmm(
             x,
@@ -411,13 +418,16 @@ def _process_tokens_locally(hidden_states_local, topk_indices_local):
             w1,
             w1_bias,
             group_sizes,
+            group_sizes_global,
             transpose_rhs=True,
             mesh=mesh,
             intermediate_size=intermediate_size,
         )
+        jax.debug.print("hidden_state after first gmm x1 {} {}", x1.sum(), x1.ravel()[:10])
+        jax.debug.print("hidden_state after first gmm x2 {} {}", x2.sum(), x2.ravel()[:10])
 
     x = activation_fn(activation, x1, x2)
-
+    jax.debug.print("hidden_state after activation {} {}", x.sum(), x.ravel()[:10])
     if use_ep:
         x = expert_sharded_gmm(
             x,
@@ -436,9 +446,11 @@ def _process_tokens_locally(hidden_states_local, topk_indices_local):
             w2,
             w2_bias,
             group_sizes,
+            group_sizes_global,
             transpose_rhs=True,
             mesh=mesh,
         )
+        jax.debug.print("hidden_state after second gmm {} {}", x.sum(), x.ravel()[:10])
 
     def _finalize_output(x_local, topk_argsort_revert_indices_local, topk_weights_local):
         x_local = x_local[topk_argsort_revert_indices_local].reshape(-1, topk, hidden_size)
@@ -453,11 +465,12 @@ def _finalize_output(x_local, topk_argsort_revert_indices_local, topk_weights_lo
         out_specs=(P(ShardingAxisName.ATTN_DATA, None)),
         check_rep=False,
     )(x, topk_argsort_revert_indices, topk_weights)
-    
+    jax.debug.print("hidden_state after finalize output {} {}", x.sum(), x.ravel()[:10])
     x = x.reshape(orig_shape)
 
     if reduce_results:
         x = jax.lax.with_sharding_constraint(x, NamedSharding(mesh, P(ShardingAxisName.ATTN_DATA)))
+        jax.debug.print("hidden_state after reducing result {} {}", x.sum(), x.ravel()[:10])
     return x