OOM on 4B model training

[armanakbari]

Hi,

I am getting OOM with 47G or GPU memory when try to train the 4B model. The neat packing is creating sequences of 40960 tokens causing the massive attention matrix that couldn't fit in memory. But even changing 40960 to 2048 gets OOM.

Any suggestions? what GPUs did you use? is there any way I can make it work with 47G memory? (I also tried ZeRO 3)

[chengshuang18]

Block diffusion training uses more memory than conventional AR models due to flex attention and 2x longer sequences. With 48GB GPU, try shorter sequences (e.g., 2048 tokens), enable gradient checkpointing, and use distributed training if possible.

[AlekseyCalvin]

But in the original issue, the OP notes that they've already tried compacting the attention range to 2048 tokens and still OOM on 47GB. I haven't yet tested your code, but as far as I know: today, in late 2025, as optimization techniques proliferate, OOM for a 4B model on 47GB VRAM and 2048 token context, at the very least, raises a lot of questions regarding the training framework implementation/recipe, even when accounting for the additional architectural demands of dLLMs and even without quantization(, etc...) accounted for.

On that point, I would highly recommend the developers, as well as the OP, to sample the dLLM quantization solutions consolidated and assessed by the QDLM (Quantization Meets dLLMs) project at https://github.com/FelixMessi/QDLM .
(Granted, there might be some bugs in the installation script for QDLM currently. So, I would also suggest you check out an attempted fix at https://github.com/hwang361/QDLM_quant ).

As it stands, with all due respect and admiration the developers, your fine-tuning guide at https://github.com/JetAstra/SDAR/tree/main/training comes across as somewhat lab-insulated, suggesting outdated expectations as to the typical early adopter user base for emerging LLM architectures, as well as the average hardware-profile/compute-budget of those attempting experimental fine-tunes of innovative open-weights models.

The wide-spread proliferation and popularization of inference and optimization solutions like oLlama, vllm, GGUF, GPTQ, bnb, and efficient fine-tuning pipelines based on TRL, Unsloth, and so forth, has gradually transformed the ecosystem as a whole. My approximation would be that, for just about any sub-40B parameter open-weights model release, out of every 10 efforts to fine tune it, 6-7 are probably non-distributed, but attempted on a single device (whether locally, or remote/rented, or via specialized environments like Colab/Modal) leveraging a single GPU, while only 2-3 might involve a distributed setup, but these are far more often than not leveraged remotely (via a rented node), than via a local hardware environment. On the other hand, local distributed setups are becoming comparatively rare, which makes a lot of sense given the much greater initial investment and complexity involved. And even labs/institutions/startups are increasingly renting compute from dedicated vendors, especially for training, rather than relying on local hardware.

[armanakbari]

Block diffusion training uses more memory than conventional AR models due to flex attention and 2x longer sequences. With 48GB GPU, try shorter sequences (e.g., 2048 tokens), enable gradient checkpointing, and use distributed training if possible.

I switched to 8 H100s and run the provided SFT code with 2048 as cutoff_len and still got OOM. I also tried deep speed Z3 but got OOM. Here is my config file:

### model
model_name_or_path: ../model/SDAR-4B-Chat
train_from_scratch: false
trust_remote_code: true

### method
stage: sft
do_train: true
finetuning_type: full
freeze_vision_tower: true
freeze_multi_modal_projector: false
freeze_language_model: true
deepspeed: examples/deepspeed/ds_z2_config.json
#examples/deepspeed/ds_z3_config.json
# disable_gradient_checkpointing: true
gradient_checkpointing: true

### dataset
dataset: open_r1_math
template: qwen3 # qwen2_vl
block_length: 4
cutoff_len: 2048
truncate_mode: drop # choices: [drop, cut] multimodal recommend drop
overwrite_cache: false
tokenized_path: ./tokenized_data
preprocessing_num_workers: 96
dataloader_num_workers: 4
neat_packing: true
# streaming: true
# max_steps: 1000

### output
output_dir: ./sdar_ckpt/sdar_4b_math_r1_cot/full/sft
logging_steps: 5
save_steps: 1206
save_total_limit: 10
plot_loss: true
overwrite_output_dir: true
save_only_model: false
report_to: tensorboard  # choices: [none, wandb, tensorboard, swanlab, mlflow]

### train
run_name: sdar_4b_math_r1_cot
include_effective_tokens_per_second: true
per_device_train_batch_size: 1 
gradient_accumulation_steps: 2 
learning_rate: 1.0e-5
num_train_epochs: 10.0
lr_scheduler_type: constant_with_warmup
warmup_ratio: 0.03
bf16: true
ddp_timeout: 180000000
resume_from_checkpoint: null

### eval
# val_size: 0.1
# per_device_eval_batch_size: 1
# eval_strategy: steps
# eval_steps: 500

And here is the error I get:

[rank2]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 12.50 GiB. GPU 2 has a total capacity of 93.00 GiB of which 3.77 GiB is free. Including non-PyTorch memory, this process has 66.78 GiB memory in use. Process 2912058 has 22.41 GiB memory in use. Of the allocated memory 43.67 GiB is allocated by PyTorch, and 20.44 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)

[chengshuang18]

I ran experiments on a single H200 (140 GB memory) with batch size = 1 and sequence length = 16 384 using the following config:

### model
model_name_or_path: /mnt/shared-storage-user/chengshuang/projects/dllm/Models/SDAR-4B-Chat
train_from_scratch: false
trust_remote_code: true

### methodD
stage: sft
do_train: true
finetuning_type: full
freeze_vision_tower: true
freeze_multi_modal_projector: false
freeze_language_model: true
deepspeed: examples/deepspeed/ds_z2_config.json
# disable_gradient_checkpointing: true
# gradient_checkpointing: true

### dataset
dataset: open_r1_math
template: qwen3 # qwen2_vl
block_length: 4
cutoff_len: 16384
truncate_mode: drop # choices: [drop, cut] 多模态推荐 drop
overwrite_cache: false
tokenized_path: /mnt/shared-storage-user/dllm-share/chengshuang/dataset/tokenized_dataset/open_r1_math_packlen_16384
preprocessing_num_workers: 96
dataloader_num_workers: 4
neat_packing: true
# streaming: true
# max_steps: 1000

### output
output_dir: /mnt/shared-storage-user/dllm-share/chengshuang/sdar_ckpt/sdar_4b_sb_math_r1_cot/full/sft
# output_dir: saves/test
logging_steps: 5
save_steps: 1206
save_total_limit: 10
plot_loss: true
overwrite_output_dir: true
save_only_model: false
report_to: tensorboard  # choices: [none, wandb, tensorboard, swanlab, mlflow]

### train
run_name: sdar_4b_sb_math_r1_cot
include_effective_tokens_per_second: true
per_device_train_batch_size: 1
gradient_accumulation_steps: 1
learning_rate: 1.0e-5
num_train_epochs: 10.0
lr_scheduler_type: constant_with_warmup
warmup_ratio: 0.03
bf16: true
ddp_timeout: 180000000
resume_from_checkpoint: null

The GPU memory usage was around 134 GB. Please note that flexattn uses dynamic compilation, so even with the same settings, memory usage can differ across GPUs. Based on my tests, an 80 GB GPU should handle sequence length = 8 192 without issues.

In your case, I suspect flexattn dynamic compilation might not be enabled, or there could be an environment configuration issue. Could you share more detailed error logs?

[armanakbari]

I could train with the proposed setting for 4B model on 4 H100s. Thanks.