OOM with Gemma3-27B-IT on DPO (8×A100, Deepspeed)

[Kshitiz-Khandel]

I'm encountering out-of-memory (OOM) errors when fine-tuning google/gemma-3-27b-it using DPO with 8×A100 (40GB) GPUs, while leveraging Deepspeed Zero3. Here's a breakdown:

1. Below is the base config file(without kernel optimisations)

base_model: google/gemma-3-27b-it
model_type: AutoModelForCausalLM
tokenizer_type: AutoTokenizer

load_in_8bit: False
load_in_4bit: True

chat_template: gemma3
rl: dpo
datasets:

• path: ./new_prompt_axolotl_train.jsonl
  type: chat_template.default
  field_messages: "prompt" ##messages
  field_chosen: "chosen"
  field_rejected: "rejected"
  message_property_mappings:
  role: role
  content: content
  roles:
  user: ["user"]
  assistant: ["assistant"]
  system: ["system"]

dataset_prepared_path:
val_set_size: 0.05
output_dir: ./outputs/dpo2/gemma3/experiment2/

sequence_len: 1024
sample_packing: false
pad_to_sequence_len: true

adapter: lora
lora_model_dir:
lora_r: 16
lora_alpha: 16
lora_dropout: 0.05
lora_target_linear: true

wandb_project:
wandb_entity:
wandb_watch:
wandb_name:
wandb_log_model:

gradient_accumulation_steps: 2
micro_batch_size: 2
num_epochs: 10
optimizer: adamw_bnb_8bit
lr_scheduler: cosine
learning_rate: 0.0002

bf16: auto
tf32: false

#gradient_checkpointing: true
resume_from_checkpoint:
logging_steps: 1
flash_attention: false
deepspeed: ./deepspeed_configs/zero3.json

warmup_steps: 10
evals_per_epoch: 4
saves_per_epoch: 1
weight_decay: 0.0
shuffle: True

2. To further optmise the training, below config options were enabled
   plugins:

• axolotl.integrations.cut_cross_entropy.CutCrossEntropyPlugin
• axolotl.integrations.liger.LigerPlugin

cut_cross_entropy: true

liger_rope: true
liger_rms_norm: true
liger_glu_activation: true
liger_layer_norm: true

gradient_checkpointing: "offload"

lora_mlp_kernel: true
lora_qkv_kernel: true
lora_o_kernel: true

However, enabling the lora_*_kernel options results in a model compatibility error:
These flags expect the model to be of type PeftModelForCausalLM, but Gemma3 is instantiated as AutoModelForCausalLM.

3. I also tried running the same setup with gemma-3-12b-it, but encountered OOM errors there as well

Any help on this would be appreciated
@NanoCode012

[NanoCode012]

Hey! I noticed you used load_in_4bit: true but adapter: lora. You may want to try with

-adapter: lora
+adapter: qlora

Let me know what happens after you set this.

[Kshitiz-Khandel]

Hey @NanoCode012. Could you please help out with the above? Thanks!

[NanoCode012]

@Kshitiz-Khandel to ensure it's not OOM, have you tried a same/smaller setup but with the smaller models?

[Kshitiz-Khandel]

For Gemma3- 4B model with the following configuration:

**Training works fine with ZeRO-1 and ZeRO-2.

Training also succeeds without using DeepSpeed.

However, with ZeRO-3, I run into out-of-memory (OOM) errors.**

Here's the current config used:

val_set_size: 0.05
output_dir: ./outputs/dpo/gemma-3/experiment3/

sequence_len: 1024
sample_packing: false
pad_to_sequence_len: true

adapter: lora
lora_model_dir:
lora_r: 16
lora_alpha: 16
lora_dropout: 0.05
lora_target_linear: true

wandb_project:
wandb_entity:
wandb_watch:
wandb_name:
wandb_log_model:

gradient_accumulation_steps: 2
micro_batch_size: 2
num_epochs: 10
optimizer: adamw_bnb_8bit
lr_scheduler: cosine
learning_rate: 0.0002

bf16: auto
tf32: false

resume_from_checkpoint:
logging_steps: 1
flash_attention: false
deepspeed: ./deepspeed_configs/zero2.json

warmup_steps: 10
evals_per_epoch: 2
saves_per_epoch: 1
weight_decay: 0.0
shuffle: true

overrides_of_model_config:
rope_scaling:
type: "linear" ## I tried reducing the context window to prevent OOM, but no help
factor: 0.5

Gemma3-12B-it
I attempted training a 12B model, but encountered OOM errors across all configurations — whether using ZeRO-1, ZeRO-2, or even without DeepSpeed.
@NanoCode012

[NanoCode012]

What's your memory usage across these runs? You've changed to use lora instead?

[Unu-tech]

@Kshitiz-Khandel have you solved this issue? I'm facing same situation.