[micahtyong]

Summary

Inside the trainer container, the Miles backend runs as one controller process and one actor process per GPU. For every forward or forward_backward call the controller builds the batch, puts it in the Ray object store, and each actor pulls the whole batch out and keeps its data-parallel share. When the actors finish, each sends its per-datum log probabilities back and the controller assembles the result.

Until now the batch crossed that boundary as nested Python lists: one int per token and one float per label weight, about 2.3 million Python objects per 128-example call for a classifier with 5 labels per position. Pickling and unpickling that many objects is slow, and it is slow by a different amount on each actor. The step-2 trace showed the cost in two places. After each actor returned, the controller spent about 0.9 s turning the returned lists into results. And at the start of every forward_backward, rank 0 sat 0.73 s at the first collective waiting for ranks 1–3 to finish unpickling. (A collective is the point where all ranks synchronize; the first one there waits for the rest.)

Arrays fix both. Numpy arrays pickle as one buffer each, Ray hands them to the actors without copying, and torch.tensor builds a GPU tensor from an array without walking it in Python. So this change keeps the rows as arrays from prepare_batch to the actors, and brings the log probabilities back as arrays.

• data.py: tokens, target_tokens, weights, advantages and sampling_logprobs are arrays. Sequence alignment and data-parallel padding pad arrays.
• runtime.py: Miles' bridge is subclassed so loss masks go in as arrays and per-datum log probabilities come back as arrays instead of through .tolist(). Weights, advantages and sampling log probabilities are sent as CPU tensors, because Miles slices those per context-parallel rank with either list + or torch.cat, and an array would reach torch.cat.
• replay_data.py: sampling-support validation accepts array target rows.

Stacks on #154, which prepared the rows with numpy but still turned them into lists at the end.

E2E results

A 4B LoRA classifier with 5 candidate labels per position on 4×H100, max_tokens_per_gpu 8192, with #157 and #152–#154 applied. Each step runs two forward calls on a frozen adapter, one forward_backward on the trained adapter, then optim_step; 128 examples of about 1,630 tokens per call. Same driver and data for both rows; 34 steps, steady-state medians.

Before WandB / After WandB

Median per step	Before	After
Step (client)	11.62 s	7.62 s
forward_backward (trainer)	5.18 s	3.68 s
Two forward calls (trainer)	4.54 s	2.61 s
Controller time after the actors return, per call	~0.9 s	< 0.01 s
Rank 0's wait at the first collective of forward_backward	0.73 s	0.00 s

Step-0 log probabilities are bit-identical before and after. The after run traced all four ranks; the wait is 0.000 s on every one of them.

🤖 Generated with Claude Code

[micahtyong]

We'll need to upstream this to miles, but otherwise lgtm