[BUG] CUDA OutOfMemoryError during backward pass in distributed training

[Bat-Reality]

Bug Description

Please provide a detailed description of the issue you encountered.

Environment Information

• Python Version: 3.12.4
• GPU: NVIDIA L20-40G * 8
• CUDA Version: 12.4
• Installation Method: git clone
• Trinity-RFT Version: 0.3.0.dev0

Steps to Reproduce

Please provide a minimal, self-contained, and reproducible example.

1. trinity run --config examples/XXX/XXX.yaml

Actual Behavior

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 5.76 GiB. 
GPU 0 has a total capacity of 44.52 GiB of which only 2–3 GiB is free. 
This process has ~41 GiB memory in use. Of the allocated memory, ~36 GiB is used by PyTorch.

ray.exceptions.RayTaskError(OutOfMemoryError): 
WorkerDict.actor_update_actor() failed with CUDA OOM

To reiterate, I am using 8 L20 GPUs with a batch size reduced to 16, max_response_tokens set to 2048, and repeat_times set to 8. However, errors still occur randomly.

Log Information

If applicable, include any relevant log output here.

[pan-x-c]

Could you provide more configuration details? Most of the examples in Trinity-RFT are based on H20-96G * 8. They need to be adjusted proportionally according to the GPU memory.

[Bat-Reality]

Could you provide more configuration details? Most of the examples in Trinity-RFT are based on H20-96G * 8. They need to be adjusted proportionally according to the GPU memory.

train.yaml

project: "grpo-single_choice"
name: "qwen3-4B-s1k"
checkpoint_root_dir: ''
algorithm:
  algorithm_type: grpo
  repeat_times: 8

model:
  model_path: ''
  max_prompt_tokens: 512
  max_response_tokens: 2048
cluster:
  node_num: 1
  gpu_per_node: 8
buffer:
  total_epochs: 8
  batch_size: 16
  explorer_input:
    taskset:
      name: gsm8k
      storage_type: file
      path: 'simplescaling/s1K-1.1'
      split: 'train'
      format:
        prompt_key: 'question'
        response_key: 'solution'
      rollout_args:
        temperature: 1.0
    eval_tasksets:
    - name: gsm8k-eval
      storage_type: file
      path: 'openai/gsm8k'
      subset_name: 'main'
      split: 'test'
      format:
        prompt_key: 'question'
        response_key: 'answer'
      workflow_args:
        use_base: false
    - name: math500
      storage_type: file
      path: 'HuggingFaceH4/MATH-500'
      split: 'test'
      format:
        prompt_key: 'problem'
        response_key: 'answer'
      workflow_args:
        use_base: false
    - name: amc23
      storage_type: file
      path: 'math-ai/amc23'
      split: 'test'
      format:
        prompt_key: 'question'
        response_key: 'answer'
      workflow_args:
        use_base: false
    - name: aime24
      storage_type: file
      path: 'math-ai/aime24'
      split: 'test'
      format:
        prompt_key: 'problem'
        response_key: 'solution'
      workflow_args:
        use_base: false
    - name: gaokao2023en
      storage_type: file
      path: 'MARIO-Math-Reasoning/Gaokao2023-Math-En'
      split: 'train'
      format:
        prompt_key: 'question'
        response_key: 'answer'
      workflow_args:
        use_base: false
    default_workflow_type: 'math_workflow'
    default_eval_workflow_type: 'math_eval_workflow'
  trainer_input:
    experience_buffer:
      name: gsm8k_multichoice_buffer
      storage_type: queue
      path: 'sqlite:///gsm8k-ginglechoice-s1k_0914.db'
explorer:
  eval_interval: 5
  runner_num: 32
  rollout_model:
    engine_type: vllm_async
    engine_num: 4
    tensor_parallel_size: 1
    enable_prefix_caching: false
    enforce_eager: true
    dtype: bfloat16
    seed: 42
synchronizer:
  sync_method: 'nccl'
  sync_interval: 5
  sync_timeout: 1200
trainer:
  trainer_type: 'verl'
  trainer_config_path: 'examples/mindgym/train_gsm8k.yaml'
  save_interval: 100

train_math.yaml

actor_rollout_ref:
  hybrid_engine: True
  model:
    external_lib: null
    override_config: { }
    enable_gradient_checkpointing: True
    use_remove_padding: True  # False
  actor:
    strategy: fsdp  # This is for backward-compatibility
    ppo_micro_batch_size_per_gpu: 4
    use_dynamic_bsz: True # False
    ppo_max_token_len_per_gpu: 22000
    grad_clip: 1.0
    ppo_epochs: 1
    shuffle: False
    ulysses_sequence_parallel_size: 1 # sp size
    optim:
      lr: 1e-6
      lr_warmup_steps: 20  # the total steps will be injected during runtime
      # min_lr_ratio: null   # only useful for warmup with cosine
      warmup_style: constant  # select from constant/cosine
      total_training_steps: -1  # must be override by program
    fsdp_config:
      wrap_policy:
        # transformer_layer_cls_to_wrap: None
        min_num_params: 0
      param_offload: False
      optimizer_offload: False
      fsdp_size: -1
  ref:
    fsdp_config:
      param_offload: False
      wrap_policy:
        # transformer_layer_cls_to_wrap: None
        min_num_params: 0
    log_prob_micro_batch_size_per_gpu: 4
    log_prob_use_dynamic_bsz: ${actor_rollout_ref.actor.use_dynamic_bsz}
    log_prob_max_token_len_per_gpu: ${actor_rollout_ref.actor.ppo_max_token_len_per_gpu}
    ulysses_sequence_parallel_size: ${actor_rollout_ref.actor.ulysses_sequence_parallel_size} # sp size

trainer:
  balance_batch: True
  # auto: find the last ckpt to resume. If can't find, start from scratch
  resume_mode: auto # or auto or resume_path if
  default_hdfs_dir: null
  remove_previous_ckpt_in_save: False
  del_local_ckpt_after_load: False
  val_before_train: False

[chenyushuo]

You can try to reduce actor_rollout_ref.actor.ppo_max_token_len_per_gpu in train_math.yaml, which can reduce the GPU memory usage.
I suggest setting it to 16000 or 10240.