Hey everyone, I have a 2x GB10 setup with 200gbe infiniband. I have been trying to utilize Megatron-Bridge (nvcr.io/nvidia/nemo:26.04.00) to leverage fp8 training and ideally better performance via transformer engine, however it looks completely unusable right now compared to HF Transformers.
Specifically, I am OOMing with extremely low batch sizes and context lengths seemingly for no reason. I have been trying to train Qwen 3.5 9B with 64 dim peft, however it looks like I am utilizing way too much vram. I thought this was raw DDP related, however it seems to be problematic when sharding my model using tensor paralleism=2. For example the config below already OOMs:
torchrun \
--nnodes="${PET_NNODES}" \
--nproc_per_node=1 \
--node_rank="${PET_NODE_RANK}" \
--master_addr="${PET_MASTER_ADDR}" \
--master_port="${PET_MASTER_PORT}" \
scripts/training/run_recipe.py \
--recipe "${MEGBRIDGE_RECIPE}" \
--dataset vlm-preloaded \
--seq_length 2048 \
--step_func qwen3_vl_step \
dataset.hf_processor_path=/workspace/megbridge/models/qwen3.5-9B-megatron/iter_0000000/tokenizer \
dataset.train_data_path=/workspace/megbridge/data/sft-smoke/training.jsonl \
dataset.valid_data_path=null \
dataset.test_data_path=null \
dataset.pack_sequences_in_batch=False \
dataset.num_workers=16 \
dataset.pin_memory=true \
train.train_iters=1000 \
train.global_batch_size=8 \
train.micro_batch_size=4 \
validation.eval_iters=0 \
validation.eval_interval=1000000 \
mixed_precision=bf16_mixed \
checkpoint.save=/workspace/megbridge/checkpoints/finetune_smoke \
checkpoint.load=null \
checkpoint.pretrained_checkpoint=/workspace/megbridge/models/qwen3.5-9B-megatron-tp2 \
checkpoint.save_interval=1000 \
model.seq_length=2048 \
model.tensor_model_parallel_size=2 \
model.pipeline_model_parallel_size=1 \
model.context_parallel_size=1 \
model.recompute_granularity=full \
model.recompute_method=uniform \
model.recompute_num_layers=1 \
model.sequence_parallel=true \
model.freeze_vision_model=true \
model.freeze_vision_projection=true \
optimizer.lr=0.0001 \
optimizer.min_lr=0.00001 \
optimizer.weight_decay=0.01 \
optimizer.optimizer=adam \
optimizer.use_precision_aware_optimizer=true \
optimizer.main_grads_dtype=bf16 \
optimizer.main_params_dtype=fp32 \
optimizer.exp_avg_dtype=bf16 \
optimizer.exp_avg_sq_dtype=bf16 \
optimizer.store_param_remainders=true \
scheduler.lr_warmup_iters=10 \
scheduler.lr_decay_style=cosine \
logger.log_interval=10 \
logger.log_params_norm=false \
logger.log_throughput=true \
logger.log_progress=true \
logger.wandb_project=test_megatron \
logger.wandb_exp_name=qwen3.5-9b-smoke \
logger.wandb_save_dir=/workspace/megbridge/runs/finetune_smoke/wandb \
rng.seed=1234 \
model.attention_backend=flash_attn \
'+peft={dim:64,alpha:64,dropout:0.0,dropout_position:pre,target_modules:[linear_qkv,linear_proj,in_proj,out_proj,linear_fc1,linear_fc2]}'
My current vram utilization with this exact script which supposedly aggressively reduces vram (e.g. via layerwise rematerialization) still reaches 110gb+ and sometimes OOMs. From what I understand, adam keeps fp32 master weights, which is fine by me, but this doesn’t explain the other lost vram.
According to raw calculations, with this exact setup:
- Global batch size 8
- full layerwise remat
- fp32 master weights
- context 2048
- Flash attention
It should be using ~47-50GB, not 220gb+.
Am I missing any critical setting, or is there some sort of bug with megatron-bridge for GB10? In this current state, megatron bridge seems completely unusable compared to raw hf transformers or unsloth.
