Shard latents_jax and prompt_embeds_jax across data_sharding for multi-host execution

Author: amepas

src/maxdiffusion/pipelines/flux/flux2klein_pipeline.py

@@ -21,10 +21,12 @@
 
 import jax
 import jax.numpy as jnp
+from jax.sharding import PartitionSpec as P
 import numpy as np
 from flax.linen import partitioning as nn_partitioning
 
 from maxdiffusion import max_logging
+from maxdiffusion.max_utils import device_put_replicated
 from ..pipeline_flax_utils import FlaxDiffusionPipeline
 from ...models.flux.transformers.transformer_flux_flax import Flux2KleinTransformer2DModel
 from ...models.vae_flax import FlaxAutoencoderKL
@@ -240,6 +242,17 @@ def __call__(
       trace["prompt_encoding"] = time.perf_counter() - t0
       max_logging.log(f" -> [TIMING] Prompt Encoding (Qwen3): {trace['prompt_encoding']:.4f} seconds ⏱️")
 
+      # Shard pipeline batch inputs across data_sharding for SPMD multi-host execution
+      data_sharding = jax.sharding.NamedSharding(self.mesh, P(*self._config.data_sharding))
+
+      def put_data_on_devices(x, sharding):
+        if hasattr(sharding, "is_fully_addressable") and sharding.is_fully_addressable:
+          return jax.device_put(x, sharding)
+        return device_put_replicated(x, sharding)
+
+      latents_jax = put_data_on_devices(latents_jax, data_sharding)
+      prompt_embeds_jax = put_data_on_devices(prompt_embeds_jax, data_sharding)
+
       # ---------------------------------------------------------------------
       # PHASE B: Denoising Loop (Flux Transformer)
       # ---------------------------------------------------------------------