Fix dp sharding for compute_logits_func (#1212)

kyuyeunk · web-flow · commit c572f98c4562 · 2025-12-02T14:11:38.000-08:00
diff --git a/tpu_inference/models/vllm/vllm_model_wrapper.py b/tpu_inference/models/vllm/vllm_model_wrapper.py
@@ -221,7 +221,7 @@ def jit_compute_logits_func(self):
         @functools.partial(
             jax.jit,
             out_shardings=(NamedSharding(self.mesh,
-                                         PartitionSpec(None, "model"))),
+                                         PartitionSpec("data", "model"))),
         )
         def compute_logits_func(
             params_and_buffers: Any,

Original file line number	Diff line number	Diff line change
`@@ -221,7 +221,7 @@ def jit_compute_logits_func(self):`
`221`	`221`	`@functools.partial(`
`222`	`222`	`jax.jit,`
`223`	`223`	`out_shardings=(NamedSharding(self.mesh,`
`224`		`- PartitionSpec(None, "model"))),`
	`224`	`+ PartitionSpec("data", "model"))),`
`225`	`225`	`)`
`226`	`226`	`def compute_logits_func(`
`227`	`227`	`params_and_buffers: Any,`