* Multi Image GRPO
* try matching trl semantics
* attn mask for multi image grpo
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix multi-image GRPO chunking and zoo guard in rl_replacements
image_sizes is now sliced on the image axis (img_start:img_end) when
the processor emits one row per image and num_images is provided;
sample-axis slicing is kept as the fallback. This restores correct
per-batch image_sizes alignment for multi-image VLM processors.
pixel_attention_mask now uses a three-way layout check: image-axis
when shape[0] matches image_grid_thw rows, pixel-row when shape[0]
matches pixel_values rows and is distinct from total_samples,
otherwise sample-axis. Prevents misalignment with image-axis grid
slicing for per-image masks and ambiguity when single-image-per-sample
shapes coincide.
cum_imgs slice indices materialize via .item to match the existing
cum_rows pattern in the same loop and avoid 0-dim tensors flowing
into a CUDA-tensor slice.
cum_rows is materialized on CPU once after construction; the
per-chunk loop uses .item on it, so keeping it on device caused a
GPU->CPU sync per iteration.
Add a one-time fail-loud guard in compute_loss when num_images is
provided but the resolved grpo_accumulated_loss source has no
num_images handling, pointing users at the corresponding unsloth_zoo
upgrade. The active GRPO path goes through grpo_accumulated_loss
(the local _get_per_token_logps and _get_per_token_logps_and_entropies
return None on the efficient path), so without this guard a stale
unsloth_zoo silently mis-slices multi-image batches.
* Tighten multi-image GRPO zoo compatibility guard
Only raise the zoo upgrade error when at least one entry in
num_images is not 1. Upstream TRL emits num_images=[1,1,...] for
any vision batch (one image per sample), and old unsloth_zoo
builds chunk those correctly because sample-axis and image-axis
slicing coincide for all-ones counts. Restricting the check to
batches with a real multi-image sample stops single-image VLM
GRPO from being needlessly broken on pre-companion zoo installs.
Prefer inspect.signature(grpo_accumulated_loss).parameters for
the num_images contract. Fall back to inspect.getsource string
matching only when the signature does not declare num_images
(e.g. the companion zoo wires it through **kwargs). The previous
try/except (TypeError, OSError) over getsource turned the guard
into a silent no-op when source files were absent; the new flow
raises in that case because the signature check will not have
proven support either.
* Consolidate multi-image GRPO chunking and zoo guard tests
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Daniel Han <danielhanchen@gmail.com>
98fde27057