MiniGPT with JAX and 3D Parallelism — Continued

Part 2: Multinode training with Jax and Ray

Part 1: Single node with Jax

Using the same hybrid-parallel setup from the previous section — combining tensor parallelism, data parallelism and pipeline parallelism — we can now scale the workload beyond a single machine.

For example, instead of running on a single g5.48xlarge instance with 8 GPUs, we can distribute the training across:

• 8 × g4dn nodes with 1 GPU each, or
• 2 × g5.12xlarge nodes with 4 GPUs each

From JAX’s perspective, these configurations can expose the same logical device mesh. The key difference is that devices are now distributed across multiple hosts rather than residing inside a single machine.

This is where globally-aware sharding becomes essential: although each process only sees its local shard, XLA must still understand how all shards together form a single global tensor spanning the entire multi-node mesh.

Ray Dataset

We first create a Distributed Ray Dataset from ‘TinyStories-1000.txt”. This dataset is aware of the Jax processes and therefore can be sharded equally among all global Jax processes across nodes. The stories are parsed just similar to the previous post in ‘single node — multiple gpus’ setting:

with open(file_path, 'r', encoding='utf-8', errors='replace') as f:
    for line in f:
        if '<|endoftext|>' in line:
            # Split on end token and process parts
            parts = line.split('<|endoftext|>')
            for i, part in enumerate(parts[:-1]):  # All but last part have end tokens
                current_story.append(part)
                story_text = ''.join(current_story).strip()
                if story_text:
                    stories.append(story_text + '<|endoftext|>')
                    if len(stories) >= max_stories:
                        break
                current_story = []

            # Last part becomes start of next story
            if parts[-1].strip():
                current_story = [parts[-1]]

            if len(stories) >= max_stories:
                break
        else:
            current_story.append(line)

We create a ray dataset from all parsed stories. Once it’s done, we pass further to map_batches to tokenize the stories. map_batches() transformations are executed in parallel across actors and tasks. This allows preprocessing steps such as tokenization or augmentation to scale independently from the training loop, utilizing cluster CPUs while JAX workers focus on GPU computation.

def load_and_preprocess_ray_data(
    file_path: str,
    batch_size: int,
    seq_length: int,
    max_stories: int = 100_000,
    shuffle: bool = False,
):
  print(f"Loading data from {file_path} (max {max_stories:,} stories)")
  # Read file in chunks to avoid loading entire file into memory
  stories = []
  current_story = []
  ds = ray.data.from_items([{"text": story} for story in stories])

  # 3. Tokenization
  def get_tokenizer():
      return tiktoken.get_encoding("gpt2")

  def tokenize(batch):
      # The tokenizer will be initialized once per worker due to `get_tokenizer` being
      # called in the worker's context when `map_batches` sets up the worker.
      tokenizer = get_tokenizer()
      input_ids = []

      for text in batch["text"]:
          tokens = tokenizer.encode(text, allowed_special={'<|endoftext|>'})

          if len(tokens) > seq_length:
              tokens = tokens[:seq_length]
          tokens.extend([0] * (seq_length - len(tokens)))
          input_ids.append(tokens)
      return {"input_ids": input_ids}

  ds = ds.map_batches(
      tokenize,
      batch_size=batch_size,
      batch_format="numpy",
  )

  # 4. Optional shuffle (distributed)
  if shuffle:
      ds = ds.random_shuffle()

  return ds

Ray JaxTrainer

Ray’s JaxTrainer v2 orchestrates distributed JAX training across multiple nodes using a DDP-style execution model.

In practice, this means Ray automatically:

• shards the dataset globally across JAX processes
• launches and manages distributed workers
• initializes the JAX distributed runtime
• and assigns process metadata such as global rank and local rank to every worker

This allows each training process to operate only on its local shard while still participating in a globally coordinated training job.

Jax process initialization for 2 workers each 4 devices

A JaxTrainer is typically configured through several key components:

• train_loop_per_workerDefines the training loop executed by each JAX worker. This function is distributed-aware and runs within the initialized JAX distributed context.
• train_loop_configContains the training configuration passed to every worker, such as model hyperparameters, batch size, sequence length, or optimizer settings.
• scaling_configSpecifies how many JAX processes (Ray workers) should be launched and what resources are allocated to each worker, such as the number of GPUs or TPUs per process.
• datasets The input dataset to be automatically sharded across workers for distributed training.
• run_configDefines runtime settings such as experiment name, persistence/storage path, runtime environment variables. Since this setup targets NVIDIA GPUs, we configure:

JAX_PLATFORM_NAME=cuda

so JAX uses the CUDA backend during execution.

config_data = load_config(config_path)["training_params"]
config = TrainingConfig(**config_data)

text_dist_dataset, estimated_batches = load_and_preprocess_ray_data(
    file_path=config.dataset_path,
    batch_size=config.global_batch_size,
    seq_length=config.seq_length,
    max_stories=config.max_stories,
    shuffle=config.shuffle_data,
)

scaling_config = ScalingConfig(num_workers=8, use_gpu=True) # Assuming 8 worker each 1 GPU
#scaling_config = ScalingConfig(
#   num_workers=2,
#   use_gpu=True,
#   resources_per_worker={"GPU": 4},
#) # Assuming 2 workers each 4 GPUs
trainer = JaxTrainer(
    train_loop_per_worker=train_loop_per_worker,
    train_loop_config={
        **dataclasses.asdict(config),
        "estimated_batches": estimated_batches,
    },
    scaling_config=scaling_config,
    datasets={"train": text_dist_dataset},
    run_config=RunConfig(
        name="minigpt_jaxtrainer_ddp",
        worker_runtime_env={
            "env_vars": {
                "JAX_PLATFORMS": "cuda",
                # Make sure to unset ``LD_LIBRARY_PATH`` if you're using CUDA devices, 
                # since ``LD_LIBRARY_PATH`` can override the CUDA libraries.
                "LD_LIBRARY_PATH": "",
            }
        },
        storage_path=config.storage_path,
    ),
)

result = trainer.fit()

The typical execution flow is as follows:

A Ray Dataset is first created on the main driver process of the Ray job and then passed into JaxTrainer through the datasets argument.

datasets={"train": train_dataset}

JaxTrainer then automatically shards the dataset across all JAX workers according to the global world size and process ranks.

An important detail is that the dataset is registered under the key "train". That same key is later referenced inside train_loop_per_worker, allowing each worker to retrieve its own shard of the distributed dataset.

On every g4dn node, Ray launches a JAX worker process executing the logic defined in:

train_loop_per_worker

Each worker operates only on its local shard of the Ray Dataset, typically accessed through:

train_ds = ray.train.get_dataset_shard("train")

This means:

• the driver initialises the distributed dataset
• Ray distributes shards across workers
• and every JAX process independently consumes only its assigned partition

Within an epoch at each node, we retrieve batch by batch via iterating the sharded dataset:

local_batch_size = config.global_batch_size / jax.process_count()

for local_batch in train_it.iter_batches(
    batch_size=local_batch_size,
    batch_format="numpy",
    prefetch_batches=2,
    drop_last=True,
):
    print(f"--- Starting training epoch {epoch} ---")

    input_ids_array = np.stack(local_batch["input_ids"])
    input_batch = jnp.array(input_ids_array)
    target_batch = prep_target_batch_2d(input_batch)
    # Modified calls to make_global_batch with adjusted global_shape
    global_x = make_global_batch(data_sharding, input_batch, (config.global_batch_size, config.seq_length))
    global_y = make_global_batch(data_sharding, target_batch, (config.global_batch_size, config.seq_length))

    train_loss = train_step(
        model,
        optimizer,
        train_metrics,
        (global_x, global_y),
        config.num_stages
    )

If the global batch_size is 128 and training runs across 8 data-parallel processes, then each process receives a local batch of:

128/8=16

So every worker operates on 16 samples locally, while collectively the system still trains on a global batch of 128.

Now comes the subtle — and most important — part of multi-node JAX/XLA training.

Each data-parallel worker already receives the correct shard of the dataset. However, the XLA compiler only sees the local tensors unless we explicitly tell it that these local arrays are actually fragments of a larger globally sharded tensor.

Without that information, the NamedSharding annotation becomes semantically incorrect because XLA assumes the array shape is only the local shape.

To fix this, we “upgrade” local batches into globally-aware arrays:

data_sharding = NamedSharding(mesh, P("batch", None))
global_x = make_global_batch(data_sharding, input_batch, (config.global_batch_size, config.seq_length))
global_y = make_global_batch(data_sharding, target_batch, (config.global_batch_size, config.seq_length))

Here:

• input_batch contains the current token sequences
• target_batch contains the expected next-token predictions
• both are still local tensors on each process

The key idea is that make_global_batch() does not gather or move data across machines. Instead, it registers metadata describing how the local tensor fits into the global distributed tensor ‘global_arr’.

with Mesh(device_mesh, axis_names=("data", "model")) as mesh:
    data_sharding = NamedSharding(mesh, P("data", None))
...
def make_global_batch(data_sharding: NamedSharding, local_arr: np.ndarray, global_shape: np.ndarray):
    # jax.make_array_from_process_local_data automatically handles the transfer 
    # from host memory (numpy) to device memory.
    global_arr = jax.make_array_from_process_local_data(data_sharding, local_arr, global_shape)
    return global_arr

A useful mental model is to think of local arrays as “cases” handled by independent sheriff’s offices, all coordinated under FBI oversight.

Jax global mesh with FBI analogy

Each sheriff’s office operates within its own jurisdiction, but its sergeants also report to FBI headquarters, which maintains a global view of the entire system.

In this analogy:

• Cases represent data shards
• Sheriff offices / regions represent data parallel groups
• Sergeants represent gpus.
• Responsibility splits between sergeants represent tensor parallelism
• the mesh is the FBI organizational map
• NamedSharding defines the assignment strategy

While each office works locally and independently, the FBI coordinates communication and keeps a global registry of:

• which cases each office is handling
• which region or office each sergeant belongs to
• and which part of the overall responsibility each sergeant is assigned.

Recall from previous post on TP+DP,

P("data", None) on input tensor (B x seq_length) means:

• shard the first dimension (the batch dimension)
• do not shard the second dimension (seq_length)

while the device mesh is 2D:

axis_names=("data", "model")

the "data" axis represents data parallelism and "model" is reserved for model/tensor parallelism.

Once the arrays are globally registered, NamedSharding finally becomes meaningful to XLA again:

with Mesh(device_mesh, axis_names=("data", "model")) as mesh:
    data_sharding = NamedSharding(mesh, P("data", None))
    model = create_model(
        rngs=nnx.Rngs(params=0, dropout=1),
        config=config,
    )

At this point, the compiler understands:

• the true global tensor shape
• how tensors are partitioned across hosts
• how collective operations should be scheduled
• and how computation maps onto the distributed device mesh

This distinction — local tensors vs globally-aware sharded tensors — is one of the most important concepts in scalable JAX training.

Recall that in the single-node setup, we could directly shard the batch using jax.device_put:

batch_to_pass = jax.device_put(
    (input_batch, target_batch),
    NamedSharding(mesh, P("batch", None)),
)
train_step(
    model,
    optimizer,
    metrics,
    batch_to_pass,
    config.num_stages,
)

In that scenario, JAX receives the full global batch and automatically partitions it across GPUs along the batch dimension.

So if the global batch size is 128 and we have 8 GPUs, JAX internally distributes the tensor into 8 shards of shape:

[16, seq_len]

Each GPU gets a different slice of the batch, while the sequence dimension remains intact. In the multi-node setup, however, the situation is slightly different. By the time data reaches the training step, JaxTrainer has already distributed the workload, meaning each process already owns its local shard of the batch. In other words, every worker starts directly with its local tensor:

[16, seq_len]

instead of the full global tensor:

[128, seq_len]

At that point, the goal is no longer to physically shard the tensor — that has already happened. Instead, we must inform XLA how these local tensors collectively form a globally sharded array.

That is precisely what make_global_batch() accomplishes.

Comparison

Since the two data-loading pipelines differ significantly — single-node uses PyGrain with SampledIndex-based sampling, while multi-node uses Ray-native data loading—it is not meaningful to directly compare epoch wall-clock time, because the input pipelines have different scheduling, prefetching behavior, and overhead profiles.

Instead, a more robust comparison is to measure:

time to reach comparable loss values and similar text-generation quality

Generated text and training loss from Ray JaxTrainer

Cost and Run-time comparison

Why multi-node is slower?

Multi-node training adds communication overhead that doesn’t exist (or is much smaller) in single-node setups:

1. Tensor Parallel (TP) communication TP requires frequent all-gather / reduce-scatter / all-reduce across GPUs. Within a node this uses fast NVLink, but across nodes it goes over much slower network (even with EFA/NCCL tuning).

→ Result: frequent synchronization stalls during each layer.

1. DDP gradient sync overhead DDP still performs gradient all-reduce every step across all replicas. Across nodes this introduces: higher latency + slower tail performance due to stragglers.

→ Result: step time becomes communication-bound instead of compute-bound.

Key takeaway

Although multi-node training increases the theoretical compute capacity, hybrid tensor parallelism (TP), data parallelism (DP) and/or pipeline parallelism (PP) setups are often communication-bound rather than compute-bound. As a result, single-node training with high intra-node bandwidth can achieve better wall-clock performance than multi-node configurations, especially when inter-node communication bandwidth is limited.

This trade-off depends heavily on the quality of the interconnect infrastructure, such as InfiniBand or other high-speed GPU networking technologies. When large single-node GPU systems are unavailable, combining multiple smaller nodes becomes a practical alternative, albeit typically at the cost of higher communication overhead and longer training times.

This demonstrates that, with Ray and JAX, multiple smaller GPU nodes can be combined to achieve training performance comparable to that of a larger node. Furthermore, by combining tensor parallelism (TP), data parallelism (DP), and pipeline parallelism (PP), it becomes possible to train models that would otherwise exceed the memory capacity of a single node.

Full python code: https://github.com/JustinDuy/mini-gpt