Skip to content

GRPO/RLOO cache_implementation="static" fails with the default fp32 model: the static cache is fp32, the keys are bf16 #7664

Open
Open
GRPO/RLOO cache_implementation="static" fails with the default fp32 model: the static cache is fp32, the keys are bf16#7664
Labels
🏋 GRPORelated to GRPO🏋 RLOORelated to RLOO🐛 bugSomething isn't working

Description

@qgallouedec
qgallouedec
opened on Oct 13, 2026

Reproduction

from datasets import load_dataset
from trl import GRPOConfig, GRPOTrainer


def reward_len(completions, **kwargs):
    return [float(len(c)) for c in completions]


trainer = GRPOTrainer(
    model="trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
    reward_funcs=reward_len,
    args=GRPOConfig(
        per_device_train_batch_size=4,
        num_generations=2,
        max_completion_length=8,
        max_steps=1,
        cache_implementation="static",
        report_to="none",
    ),
    train_dataset=load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train"),
)
trainer.train()

outputs (TRL main, on a Trainium2 instance; the same error on an H100: index_copy_(): self and source expected to have the same dtype, but got (self) Float and (source) BFloat16):

Traceback (most recent call last):
  File "/home/qgallouedec/issues/06-static-cache-dtype/repro.py", line 22, in <module>
    trainer.train()
    ~~~~~~~~~~~~~^^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/transformers/trainer.py", line 1562, in train
    return inner_training_loop(
        args=args,
    ...<2 lines>...
        ignore_keys_for_eval=ignore_keys_for_eval,
    )
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/transformers/trainer.py", line 1646, in _inner_training_loop
    self._run_epoch(
    ~~~~~~~~~~~~~~~^
        model=model,
        ^^^^^^^^^^^^
    ...<9 lines>...
        steps_trained_in_current_epoch=steps_trained_in_current_epoch,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/transformers/trainer.py", line 1892, in _run_epoch
    tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
  File "/tmp/trl-main/trl/trainer/grpo_trainer.py", line 1473, in training_step
    output = super().training_step(model, inputs, num_items_in_batch)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/transformers/trainer.py", line 2059, in training_step
    inputs = self._prepare_inputs(inputs)
  File "/tmp/trl-main/trl/extras/profiling.py", line 211, in wrapper
    return func(self, *args, **kwargs)
  File "/tmp/trl-main/trl/trainer/grpo_trainer.py", line 1502, in _prepare_inputs
    generation_batch = self._generate_and_score_completions(generation_batch)
  File "/tmp/trl-main/trl/trainer/grpo_trainer.py", line 2365, in _generate_and_score_completions
    ) = self._generate(prompts)
        ~~~~~~~~~~~~~~^^^^^^^^^
  File "/tmp/trl-main/trl/trainer/grpo_trainer.py", line 2157, in _generate
    completion_ids, logprobs = self._generate_single_turn(
                               ~~~~~~~~~~~~~~~~~~~~~~~~~~^
        prompt_ids, images, multimodal_fields, num_generations
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/tmp/trl-main/trl/trainer/grpo_trainer.py", line 1815, in _generate_single_turn
    prompt_completion_ids = unwrapped_model.generate(
        **generate_inputs, generation_config=self.generation_config
    )
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/utils/_contextlib.py", line 124, in decorate_context
    return func(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/transformers/generation/utils.py", line 2982, in generate
    result = decoding_method(
        self,
    ...<5 lines>...
        **model_kwargs,
    )
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/transformers/generation/utils.py", line 3182, in _sample
    outputs = self._prefill(
        input_ids,
    ...<2 lines>...
        is_first_iteration=not generation_config.is_assistant,
    )
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/transformers/generation/utils.py", line 4293, in _prefill
    return self(**model_inputs, return_dict=True)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/nn/modules/module.py", line 1783, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/nn/modules/module.py", line 1794, in _call_impl
    return forward_call(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/accelerate/utils/operations.py", line 943, in forward
    return model_forward(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/accelerate/utils/operations.py", line 931, in __call__
    return convert_to_fp32(self.model_forward(*args, **kwargs))
                           ~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/amp/autocast_mode.py", line 44, in decorate_autocast
    return func(*args, **kwargs)
  File "/tmp/trl-main/trl/trainer/utils.py", line 1611, in _fused_forward
    return original_forward(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/transformers/utils/generic.py", line 943, in wrapper
    output = func(self, *args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/transformers/models/qwen2/modeling_qwen2.py", line 444, in forward
    outputs: BaseModelOutputWithPast = self.model(
                                       ~~~~~~~~~~^
        input_ids=input_ids,
        ^^^^^^^^^^^^^^^^^^^^
    ...<5 lines>...
        **kwargs,
        ^^^^^^^^^
    )
    ^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/nn/modules/module.py", line 1783, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/nn/modules/module.py", line 1794, in _call_impl
    return forward_call(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/transformers/utils/generic.py", line 1081, in wrapper
    output = func(self, *args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/transformers/utils/output_capturing.py", line 287, in wrapper
    outputs = func(self, *args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/transformers/models/qwen2/modeling_qwen2.py", line 380, in forward
    hidden_states = decoder_layer(
        hidden_states,
    ...<5 lines>...
        **kwargs,
    )
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/transformers/modeling_layers.py", line 110, in __call__
    return super().__call__(*args, **kwargs)
           ~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/nn/modules/module.py", line 1783, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/nn/modules/module.py", line 1794, in _call_impl
    return forward_call(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/transformers/models/qwen2/modeling_qwen2.py", line 274, in forward
    hidden_states, _ = self.self_attn(
                       ~~~~~~~~~~~~~~^
        hidden_states=hidden_states,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    ...<5 lines>...
        **kwargs,
        ^^^^^^^^^
    )
    ^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/nn/modules/module.py", line 1783, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/nn/modules/module.py", line 1794, in _call_impl
    return forward_call(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/transformers/models/qwen2/modeling_qwen2.py", line 206, in forward
    key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
                               ~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/transformers/cache_utils.py", line 1440, in update
    keys, values = self.layers[layer_idx].update(key_states, value_states, *args, **kwargs)
                   ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/transformers/cache_utils.py", line 493, in update
    self.values.index_copy_(2, cache_position, value_states)
    ~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/_dynamo/eval_frame.py", line 1548, in _fn
    return fn(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch_neuronx/python_ops/torch_mlir/kernel.py", line 1788, in lower
    ) = self._get_or_compile_hlo(
        ~~~~~~~~~~~~~~~~~~~~~~~~^
        processed_inputs,
        ^^^^^^^^^^^^^^^^^
    ...<3 lines>...
        use_complex_decomposition=use_complex_decomposition,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch_neuronx/python_ops/torch_mlir/kernel.py", line 1411, in _get_or_compile_hlo
    entry = self._generate_hlo(
        inputs,
    ...<2 lines>...
        use_complex_decomposition=use_complex_decomposition,
    )
  File "/home/qgallouedec/.local/share/uv/python/cpython-3.14.3-linux-x86_64-gnu/lib/python3.14/contextlib.py", line 85, in inner
    return func(*args, **kwds)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch_neuronx/python_ops/torch_mlir/kernel.py", line 712, in _generate_hlo
    decomposed_gm = self._generate_fx_graph(
        fake_inputs,
    ...<2 lines>...
        use_complex_decomposition=use_complex_decomposition,
    )
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch_neuronx/python_ops/torch_mlir/kernel.py", line 614, in _generate_fx_graph
    decomposed_gm = make_fx(functionalized_fn, decomposition_table=decomposition_table)(
        *flat_args
    )
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/fx/experimental/proxy_tensor.py", line 3422, in wrapped
    return make_fx_tracer.trace(f, *args)
           ~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/fx/experimental/proxy_tensor.py", line 3298, in trace
    return self._trace_inner(f, *args)
           ~~~~~~~~~~~~~~~~~^^^^^^^^^^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/fx/experimental/proxy_tensor.py", line 3248, in _trace_inner
    t = dispatch_trace(
        wrap_key(func, args, self.fx_tracer, self.pre_dispatch),
        tracer=self.fx_tracer,
        concrete_args=tuple(phs),
    )
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/_compile.py", line 54, in inner
    return disable_fn(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/_dynamo/eval_frame.py", line 1548, in _fn
    return fn(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/fx/experimental/proxy_tensor.py", line 1863, in dispatch_trace
    graph = tracer.trace(root, concrete_args)  # type: ignore[arg-type]
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/fx/_symbolic_trace.py", line 932, in trace
    (self.create_arg(fn(*args)),),
                     ~~^^^^^^^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/fx/experimental/proxy_tensor.py", line 1952, in wrapped
    out = f(*tensors)  # type:ignore[call-arg]
  File "<string>", line 1, in <lambda>
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch_neuronx/python_ops/torch_mlir/kernel.py", line 588, in functionalized_fn
    result = wrapped_fn(*func_args)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch_neuronx/python_ops/torch_mlir/kernel.py", line 563, in wrapped_fn
    result = flat_fn(*args)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/_functorch/_aot_autograd/utils.py", line 217, in flat_fn
    tree_out = fn(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch_neuronx/python_ops/torch_mlir/kernel.py", line 224, in wrapped
    result = fn(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch_neuronx/python_ops/torch_mlir/ops/indexing.py", line 247, in torch_index_copy
    return torch.index_copy(x, dim, indexes, source)
           ~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/fx/experimental/proxy_tensor.py", line 2014, in __torch_function__
    return func(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/_compile.py", line 54, in inner
    return disable_fn(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/_dynamo/eval_frame.py", line 1548, in _fn
    return fn(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/_subclasses/functional_tensor.py", line 638, in __torch_dispatch__
    outs_unwrapped = func._op_dk(
        torch._C.DispatchKey.Functionalize,
        *args_unwrapped,
        **kwargs_unwrapped,
    )
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/_compile.py", line 54, in inner
    return disable_fn(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/_dynamo/eval_frame.py", line 1548, in _fn
    return fn(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/utils/_stats.py", line 29, in wrapper
    return fn(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/fx/experimental/proxy_tensor.py", line 2198, in __torch_dispatch__
    return proxy_call(self, func, self.pre_dispatch, args, kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/fx/experimental/proxy_tensor.py", line 1328, in proxy_call
    r = maybe_handle_decomp(proxy_mode, func, args, kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/fx/experimental/proxy_tensor.py", line 3495, in maybe_handle_decomp
    out = decomp_table[op](*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch_neuronx/neuron_dynamo_backend/decompositions.py", line 2998, in index_copy
    return torch.scatter(input, dim, index_expanded, source)
           ~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/_compile.py", line 54, in inner
    return disable_fn(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/_dynamo/eval_frame.py", line 1548, in _fn
    return fn(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/utils/_stats.py", line 29, in wrapper
    return fn(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/fx/experimental/proxy_tensor.py", line 2198, in __torch_dispatch__
    return proxy_call(self, func, self.pre_dispatch, args, kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/fx/experimental/proxy_tensor.py", line 1449, in proxy_call
    out = func(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/_ops.py", line 916, in __call__
    return self._op(*args, **kwargs)
           ~~~~~~~~^^^^^^^^^^^^^^^^^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/_prims_common/wrappers.py", line 313, in _fn
    result = fn(*args, **kwargs)
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/_meta_registrations.py", line 6249, in meta_scatter
    scatter_meta_impl(self, dim, index, src, reduce)
    ~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/_meta_registrations.py", line 6219, in scatter_meta_impl
    scatter_gather_dtype_check("scatter", self, index, src)
    ~~~~~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/_meta_registrations.py", line 6165, in scatter_gather_dtype_check
    torch._check(
    ~~~~~~~~~~~~^
        self.dtype == src_opt.dtype,
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
        lambda: f"{method_name}(): Expected self.dtype to be equal to src.dtype",
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    )
    ^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/__init__.py", line 2140, in _check
    _check_with(RuntimeError, cond, message)
    ~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/venvs/neuron-main/lib/python3.14/site-packages/torch/__init__.py", line 2118, in _check_with
    raise error_type(message_evaluated)
RuntimeError: scatter(): Expected self.dtype to be equal to src.dtype

TRL loads the model in fp32 by default (model_init_kwargs dtype defaults to "float32"), and bf16=True is the default, so accelerate wraps the model's forward in bf16 autocast. Generation goes through that forward and produces bf16 key/value states. StaticCache allocates its buffers in the model dtype (fp32), and update writes the bf16 states into them with index_copy_, which requires matching dtypes. With the default dynamic cache nothing breaks, because it stores whatever dtype it receives.

Loading the model in bf16 avoids it (model_init_kwargs={"dtype": "bfloat16"}): GRPO then trains with the static cache on H100, and GRPO and RLOO do on Trainium.

How did you hit this?

Testing GRPO on Trainium, where a static cache matters more than anywhere else: with the dynamic cache, every new KV length is a new shape, and each new shape is a recompilation. GRPO on Qwen3-0.6B with the dynamic cache takes 959 s for step 1, 222 s for step 2 and over 12 min for step 3, see #7658. Not Neuron-specific: same failure on CUDA with the same script.

System Info

- Platform: Linux-6.12.94-123.192.amzn2023.x86_64-x86_64-with-glibc2.39
- Python version: 3.14.3
- TRL version: 1.16.0.dev0+07baa40
- PyTorch version: 2.14.0+cpu
- torch-neuronx version: 2.14.3.0.2115 (TorchNeuron, native backend), neuronx-cc 2.28.4568.0
- accelerator(s): cpu (as printed; the machine is a trn2 instance with 1 Trainium2 device, 4 logical NeuronCores, see #7663)
- Transformers version: 5.19.0
- Accelerate version: 1.15.0
- Accelerate config: not found
- Datasets version: 5.1.0
- HF Hub version: 2.2.0
- bitsandbytes version: not installed
- DeepSpeed version: not installed
- Liger-Kernel version: not installed
- PEFT version: 0.21.2
- vLLM version: not installed

Activity

  1. qgallouedec
    added
    🐛 bugSomething isn't working
    🏋 RLOORelated to RLOO
    🏋 GRPORelated to GRPO
    on Oct 13, 2026
  2. qgallouedec
    added a parent issue on Oct 13, 2026

  3. qgallouedec commented on Oct 13, 2026

    @qgallouedec
    qgallouedec
    on Oct 13, 2026
    MemberAuthor

    Correction: the dtype mismatch is between keys and values, not cache vs model. Under bf16 autocast with an fp32 model, the keys reach the cache in fp32 (RoPE multiplies them by fp32 cos/sin) while the values are bf16, and StaticLayer.lazy_initialization allocates both buffers with key_states.dtype. So the fix belongs in transformers: allocate the values with value_states.dtype. PR incoming there.

  4. qgallouedec commented on Oct 13, 2026

    @qgallouedec
    qgallouedec
    on Oct 13, 2026
    MemberAuthor

    Fix: huggingface/transformers#49554

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    🏋 GRPORelated to GRPO🏋 RLOORelated to RLOO🐛 bugSomething isn't working

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions