Repository: unslothai/unsloth
File: unsloth/models/rl_replacements.py (Python)

Code before:
                max_left_pad = 0

            input_ids_chunks = torch.chunk(input_ids, chunks = B, dim = 0)
            attention_mask_chunks = torch.chunk(attention_mask, chunks = B, dim = 0)

            def chunk_optional(tensor, chunks):
                if tensor is None:
                    return [None] * chunks
                return torch.chunk(tensor, chunks = chunks, dim = 0)

            pixel_values_chunks = [None] * B
            image_grid_thw_chunks = [None] * B
            pixel_attention_mask_chunks = [None] * B

            # This is the chunkng logit from trl 0.23.0
            if image_grid_thw is not None and pixel_values is not None:
                if image_grid_thw.shape[0] != B:
                    raise ValueError(
                        f"This logic requires image_grid_thw.shape[0] ({image_grid_thw.shape[0]}) "
                        f"to be equal to batch size B ({B})."
                    )

                rows_per_sample = image_grid_thw.prod(dim = -1)
                rows_per_sample_list = rows_per_sample.cpu().tolist()

                pixel_values_chunks = list(
                    torch.split(pixel_values, rows_per_sample_list, dim = 0)
                )
                if pixel_attention_mask is not None:
                    pixel_attention_mask_chunks = list(
                        torch.split(pixel_attention_mask, rows_per_sample_list, dim = 0)
                    )

                image_grid_thw_chunks = list(
                    torch.chunk(image_grid_thw, chunks = B, dim = 0)
                )

            elif pixel_values is not None:
                pixel_values_chunks = list(torch.chunk(pixel_values, chunks = B, dim = 0))
                if pixel_attention_mask is not None:
                    pixel_attention_mask_chunks = list(
                        torch.chunk(pixel_attention_mask, chunks = B, dim = 0)
                    )

            if image_sizes is not None and not isinstance(image_sizes, torch.Tensor):
                image_sizes_chunks = [[size] for size in image_sizes]
            else:
                image_sizes_chunks = chunk_optional(image_sizes, B)

            lm_head = self.model.get_output_embeddings().weight
            temperature = self.temperature

Review comment (suggestion):
This list is only to ensure python iterability I suppose?
Can you please run on some dummy inputs that this (just split-list of tensor as standalone and nothing else) doesn't replicate the data? Ideally it should not but good to validate.

Code after:
        compute_efficient = False,
        *args,
        **kwargs,
    ):
        # All Unsloth code here in this function is licensed under AGPL3
        # if True: # os.environ.get('UNSLOTH_USE_NEW_MODEL', '0') == '0':
        #     return None, None  # logps, entropies Unsloth efficient GRPO
        if compute_efficient:
            return None, None
        else:
            if not hasattr(self, "_autocast_dtype"):
                self._autocast_dtype = (
                    torch.float16
                    if os.environ.get("ACCELERATE_MIXED_PRECISION", "fp16") == "fp16"
                    else torch.bfloat16
                )
                if os.environ.get("UNSLOTH_FORCE_FLOAT32", "0") == "1":
                    self._autocast_dtype = torch.float16

            pixel_values, image_grid_thw = (
                kwargs.get("pixel_values", None),
                kwargs.get("image_grid_thw", None),
            )
            pixel_attention_mask, image_sizes = (
                kwargs.get("pixel_attention_mask", None),
                kwargs.get("image_sizes", None),
            )

            unwrapped_model = self.accelerator.unwrap_model(
                model, keep_fp32_wrapper = False
            )

            lm_head = self.model.get_output_embeddings().weight

            dtype_bytes = (
                16 if self._autocast_dtype in [torch.float16, torch.bfloat16] else 32
            )
            total_rows = input_ids.shape[0]
            seq_len = input_ids.shape[1]
            hidden_dim = lm_head.shape[1]
            vocab_dim = lm_head.shape[0]

            if self.args.unsloth_grpo_mini_batch is None:
                B, multiplier = autotune_batch_and_chunks(
                    total_rows,
                    seq_len,
                    hidden_dim,
                    vocab_dim,
                    dtype_bytes,
                    self.args.unsloth_logit_chunk_multiplier,
                )