Error when running GRPO in examples/refine.py: ValueError: `images` are expected as arguments to a `DetikzifyProcessor` instance.

[MqLeet]

Hi @potamides, thanks for your great work!

I met the error below when I run the code in examples/refine.py:

[rank0]: Traceback (most recent call last):
[rank0]:   File "/cpfs04/user/maqianli/DeTikZify/examples/refine.py", line 338, in <module>
[rank0]:     train(model=model, processor=processor, dataset=dataset, **args)
[rank0]:   File "/cpfs04/user/maqianli/DeTikZify/examples/refine.py", line 254, in train
[rank0]:     trainer.train(resume_from_checkpoint=None if overwrite else get_last_checkpoint(output_dir))
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/transformers/trainer.py", line 2240, in train
[rank0]:     return inner_training_loop(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/transformers/trainer.py", line 2555, in _inner_training_loop
[rank0]:     tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
[rank0]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/transformers/trainer.py", line 3739, in training_step
[rank0]:     inputs = self._prepare_inputs(inputs)
[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/trl/extras/profiling.py", line 98, in wrapper
[rank0]:     return func(self, *args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/trl/trainer/grpo_trainer.py", line 989, in _prepare_inputs
[rank0]:     generation_batch = self._generate_and_score_completions(generation_batch)
[rank0]:                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/trl/trainer/grpo_trainer.py", line 1058, in _generate_and_score_completions
[rank0]:     prompt_inputs = self.processing_class(
[rank0]:                     ^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/DeTikZify/detikzify/model/processing_detikzify.py", line 93, in __call__
[rank0]:     raise ValueError("`images` are expected as arguments to a `DetikzifyProcessor` instance.")
[rank0]: ValueError: `images` are expected as arguments to a `DetikzifyProcessor` instance.

My bash scripts is:

export CUDA_VISIBLE_DEVICES=0 
export WORLD_SIZE=1


torchrun --nproc_per_node=1 examples/refine.py \
    --base_model "models/nllg/detikzify-v2.5-8b" \
    --datikz "nllg/datikz-v3" \
    --output "outputs/detikzify-v2.5-8b-posttrain-single" \
    --batch_size 1 \
    --num_train_steps 100 \
    --sketch_ratio 0.0 \
    --num_compile_workers 2 \
    --gradient_checkpointing

How can I solve this?

[potamides]

Hi @MqLeet, thank you for your interest in this project :)

How did you install trl? As of now, GRPO with image support is not yet merged into the main branch, so you have to install it from a pull request and apply some patches for it to work:

git clone https://github.com/hellopahe/trl.git
curl -L http://github.com/huggingface/trl/pull/3568.patch | git -C trl apply
pip install trl

Note that these instructions would also be printed if you try to invoke the script without installing trl first:

DeTikZify/examples/refine.py

Lines 36 to 46 in 24bb0ef

	try:
	from trl import GRPOConfig, GRPOTrainer
	except ModuleNotFoundError:
	print(
	"You need to install trl with vision support to be able to use this script:",
	"git clone https://github.com/hellopahe/trl.git",
	"curl -L http://github.com/huggingface/trl/pull/3568.patch | git -C trl apply",
	"pip install trl",
	sep="\n\t"
	)
	exit(1)

[MqLeet]

Thank you for your immediate reply! But I meet a new bug: shape mismatch 😭 How should I solve this?

[rank0]: Traceback (most recent call last):
[rank0]:   File "/cpfs04/user/maqianli/DeTikZify/examples/refine.py", line 338, in <module>
[rank0]:     train(model=model, processor=processor, dataset=dataset, **args)
[rank0]:   File "/cpfs04/user/maqianli/DeTikZify/examples/refine.py", line 254, in train
[rank0]:     trainer.train(resume_from_checkpoint=None if overwrite else get_last_checkpoint(output_dir))
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/transformers/trainer.py", line 2240, in train
[rank0]:     return inner_training_loop(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/transformers/trainer.py", line 2555, in _inner_training_loop
[rank0]:     tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
[rank0]:                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/transformers/trainer.py", line 3739, in training_step
[rank0]:     inputs = self._prepare_inputs(inputs)
[rank0]:              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/DeTikZify/trl/trl/extras/profiling.py", line 96, in wrapper
[rank0]:     return func(self, *args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/DeTikZify/trl/trl/trainer/grpo_trainer.py", line 1019, in _prepare_inputs
[rank0]:     generation_batch = self._generate_and_score_completions(generation_batch)
[rank0]:                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/DeTikZify/trl/trl/trainer/grpo_trainer.py", line 1209, in _generate_and_score_completions
[rank0]:     prompt_completion_ids = unwrapped_model.generate(
[rank0]:                             ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/transformers/generation/utils.py", line 2597, in generate
[rank0]:     result = self._sample(
[rank0]:              ^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/transformers/generation/utils.py", line 3560, in _sample
[rank0]:     outputs = model_forward(**model_inputs, return_dict=True)
[rank0]:               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/accelerate/utils/operations.py", line 818, in forward
[rank0]:     return model_forward(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/accelerate/utils/operations.py", line 806, in __call__
[rank0]:     return convert_to_fp32(self.model_forward(*args, **kwargs))
[rank0]:                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/torch/amp/autocast_mode.py", line 44, in decorate_autocast
[rank0]:     return func(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/DeTikZify/detikzify/model/modeling_detikzify.py", line 343, in forward
[rank0]:     outputs = self.model(
[rank0]:               ^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1739, in _wrapped_call_impl
[rank0]:     return self._call_impl(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1750, in _call_impl
[rank0]:     return forward_call(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/DeTikZify/detikzify/model/modeling_detikzify.py", line 245, in forward
[rank0]:     inputs_embeds = self.inputs_merger(
[rank0]:                     ^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/DeTikZify/detikzify/model/modeling_detikzify.py", line 178, in inputs_merger
[rank0]:     new_inputs_embeds[special_image_token_mask] = reshaped_image_hidden_states
[rank0]:     ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]: RuntimeError: shape mismatch: value tensor of shape [9600, 4096] cannot be broadcast to indexing result of shape [0, 4096]

[potamides]

I think this could be related to gradient checkpointing, could you try again without --gradient_checkpointing?

[MqLeet]

I think this could be related to gradient checkpointing, could you try again without --gradient_checkpointing?

Hi @potamides, thanks for your quick reply! I remove the --gradient_checkpointing, but the line 178 in detikzify/model/modeling_detikzify.py still get an error:

INFO|trainer.py:706] 2025-07-09 11:35:22,804 >> max_steps is given, it will override any value given in num_train_epochs
[INFO|trainer.py:756] 2025-07-09 11:35:22,805 >> Using auto half precision backend
[INFO|trainer.py:2409] 2025-07-09 11:35:24,487 >> ***** Running training *****
[INFO|trainer.py:2410] 2025-07-09 11:35:24,488 >>   Num examples = 50
[INFO|trainer.py:2411] 2025-07-09 11:35:24,488 >>   Num Epochs = 1
[INFO|trainer.py:2412] 2025-07-09 11:35:24,488 >>   Instantaneous batch size per device = 1
[INFO|trainer.py:2415] 2025-07-09 11:35:24,488 >>   Total train batch size (w. parallel, distributed & accumulation) = 32
[INFO|trainer.py:2416] 2025-07-09 11:35:24,488 >>   Gradient Accumulation steps = 32
[INFO|trainer.py:2417] 2025-07-09 11:35:24,488 >>   Total optimization steps = 50
[INFO|trainer.py:2418] 2025-07-09 11:35:24,490 >>   Number of trainable parameters = 8,044,417,024
  0%|                                                                                                                      | 0/50 [00:00<?, ?it/s]/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/_distutils_hack/__init__.py:53: UserWarning: Reliance on distutils from stdlib is deprecated. Users must rely on setuptools to provide the distutils module. Avoid importing distutils or import setuptools first, and avoid setting SETUPTOOLS_USE_DISTUTILS=stdlib. Register concerns at https://github.com/pypa/setuptools/issues/new?template=distutils-deprecation.yml
  warnings.warn(
/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/_distutils_hack/__init__.py:53: UserWarning: Reliance on distutils from stdlib is deprecated. Users must rely on setuptools to provide the distutils module. Avoid importing distutils or import setuptools first, and avoid setting SETUPTOOLS_USE_DISTUTILS=stdlib. Register concerns at https://github.com/pypa/setuptools/issues/new?template=distutils-deprecation.yml
  warnings.warn(
/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/_distutils_hack/__init__.py:53: UserWarning: Reliance on distutils from stdlib is deprecated. Users must rely on setuptools to provide the distutils module. Avoid importing distutils or import setuptools first, and avoid setting SETUPTOOLS_USE_DISTUTILS=stdlib. Register concerns at https://github.com/pypa/setuptools/issues/new?template=distutils-deprecation.yml
  warnings.warn(
/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/_distutils_hack/__init__.py:53: UserWarning: Reliance on distutils from stdlib is deprecated. Users must rely on setuptools to provide the distutils module. Avoid importing distutils or import setuptools first, and avoid setting SETUPTOOLS_USE_DISTUTILS=stdlib. Register concerns at https://github.com/pypa/setuptools/issues/new?template=distutils-deprecation.yml
  warnings.warn(
/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/_distutils_hack/__init__.py:53: UserWarning: Reliance on distutils from stdlib is deprecated. Users must rely on setuptools to provide the distutils module. Avoid importing distutils or import setuptools first, and avoid setting SETUPTOOLS_USE_DISTUTILS=stdlib. Register concerns at https://github.com/pypa/setuptools/issues/new?template=distutils-deprecation.yml
  warnings.warn(
[WARNING|logging.py:328] 2025-07-09 11:35:27,500 >> `generation_config` default values have been modified to match model-specific defaults: {'max_length': 4096, 'temperature': 0.8}. If this is not desired, please set these values explicitly.

[rank0]:W0709 11:38:09.080000 1284035 site-packages/torch/_logging/_internal.py:1089] [0/0] Profiler function <class 'torch.autograd.profiler.record_function'> will be ignored
/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/torch/autograd/graph.py:823: UserWarning: Error detected in IndexPutBackward0. Traceback of forward call that caused the error:
  File "/cpfs04/user/maqianli/DeTikZify/detikzify/model/modeling_detikzify.py", line 247, in torch_dynamo_resume_in_forward_at_235
    inputs_embeds = self.inputs_merger(
  File "/cpfs04/user/maqianli/DeTikZify/detikzify/model/modeling_detikzify.py", line 180, in inputs_merger
    new_inputs_embeds[special_image_token_mask] = reshaped_image_hidden_states
 (Triggered internally at /pytorch/torch/csrc/autograd/python_anomaly_mode.cpp:122.)
  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
[rank0]:W0709 11:39:02.014000 1284035 site-packages/torch/_dynamo/exc.py:304] [13/0] Backend compiler failed with a fake tensor exception at 
[rank0]:W0709 11:39:02.014000 1284035 site-packages/torch/_dynamo/exc.py:304] [13/0]   File "/cpfs04/user/maqianli/DeTikZify/detikzify/model/modeling_detikzify.py", line 267, in torch_dynamo_resume_in_forward_at_235
[rank0]:W0709 11:39:02.014000 1284035 site-packages/torch/_dynamo/exc.py:304] [13/0]     return DetikzifyBaseModelOutputWithPast(
[rank0]:W0709 11:39:02.014000 1284035 site-packages/torch/_dynamo/exc.py:304] [13/0] Adding a graph break.
/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/torch/autograd/graph.py:823: UserWarning: Error detected in IndexPutBackward0. Traceback of forward call that caused the error:
  File "/cpfs04/user/maqianli/DeTikZify/detikzify/model/modeling_detikzify.py", line 247, in torch_dynamo_resume_in_forward_at_235
    inputs_embeds = self.inputs_merger(
  File "/cpfs04/user/maqianli/DeTikZify/detikzify/model/modeling_detikzify.py", line 180, in inputs_merger
    new_inputs_embeds[special_image_token_mask] = reshaped_image_hidden_states
 (Triggered internally at /pytorch/torch/csrc/autograd/python_anomaly_mode.cpp:122.)
  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
[rank0]:W0709 11:39:08.386000 1284035 site-packages/torch/_dynamo/exc.py:304] [13/0_1] Backend compiler failed with a fake tensor exception at 
[rank0]:W0709 11:39:08.386000 1284035 site-packages/torch/_dynamo/exc.py:304] [13/0_1]   File "/cpfs04/user/maqianli/DeTikZify/detikzify/model/modeling_detikzify.py", line 267, in torch_dynamo_resume_in_forward_at_235
[rank0]:W0709 11:39:08.386000 1284035 site-packages/torch/_dynamo/exc.py:304] [13/0_1]     return DetikzifyBaseModelOutputWithPast(
[rank0]:W0709 11:39:08.386000 1284035 site-packages/torch/_dynamo/exc.py:304] [13/0_1] Adding a graph break.
/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/torch/autograd/graph.py:823: UserWarning: Error detected in IndexPutBackward0. Traceback of forward call that caused the error:
  File "/cpfs04/user/maqianli/DeTikZify/detikzify/model/modeling_detikzify.py", line 180, in inputs_merger
    new_inputs_embeds[special_image_token_mask] = reshaped_image_hidden_states
 (Triggered internally at /pytorch/torch/csrc/autograd/python_anomaly_mode.cpp:122.)
  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
[rank0]:W0709 11:39:08.678000 1284035 site-packages/torch/_dynamo/exc.py:304] [15/0] Backend compiler failed with a fake tensor exception at 
[rank0]:W0709 11:39:08.678000 1284035 site-packages/torch/_dynamo/exc.py:304] [15/0]   File "/cpfs04/user/maqianli/DeTikZify/detikzify/model/modeling_detikzify.py", line 181, in inputs_merger
[rank0]:W0709 11:39:08.678000 1284035 site-packages/torch/_dynamo/exc.py:304] [15/0]     return new_inputs_embeds
[rank0]:W0709 11:39:08.678000 1284035 site-packages/torch/_dynamo/exc.py:304] [15/0] Adding a graph break.
/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/torch/autograd/graph.py:823: UserWarning: Error detected in IndexPutBackward0. Traceback of forward call that caused the error:
  File "/cpfs04/user/maqianli/DeTikZify/detikzify/model/modeling_detikzify.py", line 180, in inputs_merger
    new_inputs_embeds[special_image_token_mask] = reshaped_image_hidden_states
 (Triggered internally at /pytorch/torch/csrc/autograd/python_anomaly_mode.cpp:122.)
  return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass

and this leads to an CUDA OOM error finally

[rank0]:W0709 11:39:08.726000 1284035 site-packages/torch/_dynamo/exc.py:304] [15/0_1] Backend compiler failed with a fake tensor exception at 
[rank0]:W0709 11:39:08.726000 1284035 site-packages/torch/_dynamo/exc.py:304] [15/0_1]   File "/cpfs04/user/maqianli/DeTikZify/detikzify/model/modeling_detikzify.py", line 181, in inputs_merger
[rank0]:W0709 11:39:08.726000 1284035 site-packages/torch/_dynamo/exc.py:304] [15/0_1]     return new_inputs_embeds
[rank0]:W0709 11:39:08.726000 1284035 site-packages/torch/_dynamo/exc.py:304] [15/0_1] Adding a graph break.
[rank0]: Traceback (most recent call last):
[rank0]:   File "/cpfs04/user/maqianli/DeTikZify/examples/refine.py", line 338, in <module>
[rank0]:     train(model=model, processor=processor, dataset=dataset, **args)
[rank0]:   File "/cpfs04/user/maqianli/DeTikZify/examples/refine.py", line 254, in train
[rank0]:     trainer.train(resume_from_checkpoint=None if overwrite else get_last_checkpoint(output_dir))
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/transformers/trainer.py", line 2240, in train
[rank0]:     return inner_training_loop(
[rank0]:            ^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/transformers/trainer.py", line 2606, in _inner_training_loop
[rank0]:     self.optimizer.step()
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/accelerate/optimizer.py", line 179, in step
[rank0]:     self.optimizer.step(closure)
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/torch/optim/lr_scheduler.py", line 140, in wrapper
[rank0]:     return func.__get__(opt, opt.__class__)(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/torch/optim/optimizer.py", line 493, in wrapper
[rank0]:     out = func(*args, **kwargs)
[rank0]:           ^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/torch/optim/optimizer.py", line 91, in _use_grad
[rank0]:     ret = func(self, *args, **kwargs)
[rank0]:           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/torch/optim/adamw.py", line 232, in step
[rank0]:     has_complex = self._init_group(
[rank0]:                   ^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 745, in _fn
[rank0]:     return fn(*args, **kwargs)
[rank0]:            ^^^^^^^^^^^^^^^^^^^
[rank0]:   File "/cpfs04/user/maqianli/miniconda3/envs/detikzify311/lib/python3.11/site-packages/torch/optim/adamw.py", line 171, in _init_group
[rank0]:     state["exp_avg"] = torch.zeros_like(
[rank0]:                        ^^^^^^^^^^^^^^^^^
[rank0]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 1002.00 MiB. GPU 0 has a total capacity of 79.35 GiB of which 932.19 MiB is free. Process 820774 has 78.43 GiB memory in use. Of the allocated memory 73.79 GiB is allocated by PyTorch, and 2.92 GiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation.  See documentation for Memory Management  (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)

I modified the codes in TrainDataset in example/refine.py, is this the cause of the error above?

    def sample(self, n_samples):
        # spiqa: Dataset = Dataset.from_generator( # type: ignore
        #     generator=self.sample_spiqa_dataset,
        #     gen_kwargs=dict(n_samples=round(.5 * n_samples))
        # )
        datikz: Dataset = Dataset.from_generator( # type: ignore
            generator=self.sample_datikz_dataset,
            gen_kwargs=dict(n_samples=round(.5 * n_samples))
        )

        return concatenate_datasets([datikz, datikz])

[potamides]

No, your modification should not be the cause. I ran into the same issue as you when using gradient checkpointing. Without it, it works for me, but my GPUs have more VRAM than yours, which is why you're experiencing CUDA OOM errors instead. Interestingly, this issue has never occurred in the past. I suspect it was introduced with more recent transformers/PyTorch versions.

The error happens in the inputs_merger function in the DeTikZify model code. Since DeTikZify's model code is based on Idefics 3, I believe we can try to fix this issue by integrating some of its recent changes.

[MqLeet]

I set the torch_compile=False, setup the deepspeed and flash_attn, then the scripts can be ran without the OOM error🧐

But the RL training is quite time-consuming, maybe 38 hours for 100 samples, is it normal?

[potamides]

Oh, cool that you found a workaround! Yeah, that looks about right to me. I trained for 500 steps on 8x H200 GPUs which took ~5 days to complete. Reducing batch size or number of completions per sample would speed up the training, of course.

[potamides]

By the way, seems like a similar bug also happens with Idefics 3, as reported in huggingface/transformers#31380 (comment). It doesn't look like there is a fix yet.