MoT_vae_stage2_instruct.yaml training slow

[DominikVincent]

Hello,

I observe that the training of this stage is much slower than the previous one. Is that expected?

For example:

1. stage1 MoT_vae_stage1_t2m.yaml the training succeeds in 2-4 batches per second
2. stage2 MoT_vae_stage2_all.yaml the training succeeds in 1.5-3 batches per second
3. stage3 MoT_vae_stage2_instruct.yaml the training succeeds in only 0.4 batches per second, together with the higher amount of batches per epoch, makes the training very slow.

Did you observe the same?

[BoifZ]

In Short

• Stage 3 can be slower because the trained text branch adds substantial forward/backward overhead.
• Stage2_instruct is typically only slightly slower than Stage 1 due to longer effective sequences.
• Your 0.4 it/s is slower than we usually observe.

What we expect

On our side (2×GPU, batch size 32, no grad accumulation), we see ~0.9 it/s across stages, with Stage 3 slightly slower than Stage 1/2 but not by 2–3×.

Why Stage2_instruct/Stage3 may be slower

• Stage2_instruct: Longer sequences (prompt + response) increase per-step compute
• Stage3: The trained text branch increases conputational overhead during forward/backward.
• If multiple jobs share the same machine, I/O/PCIe/CPU contention can further reduce throughput.

Check configures to reproduce the results

• Hardware: GPU model and count (e.g., 2×3090 24GB), DEVICE, BATCH_SIZE.
• Versions: PyTorch / CUDA / Lightning; torch.backends.cudnn.deterministic.
• Precision: AMP (fp16/bf16) status, and accumulate_grad_batches .
• Dataloader: num_workers, pin_memory, persistent_workers, and storage (NVMe vs HDD/NFS).
• Concurrency: other training jobs on the same node.

Our reference settings & logs

Global

DEVICE: [0,1]
BATCH_SIZE: 32
NUM_WORKERS: 8
TRAIN.accumulate_grad_batches: 1

Stage1 (STAGE: lm_pretrain, instruction_type: t2m)

Epoch 0/299 ━━━━━━━━━━━━━━╸━━━━━━━━━━━━━━━━━━━━━━━━━ 145/400 0:02:39 • 0:04:39 0.92it/s 

Stage2 (STAGE: lm_pretrain, instruction_type: all)

Epoch 0/299 ━━━━━━━━━━╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 303/1200 0:05:28 • 0:16:16 0.92it/s  

Stage2_instruct_all (STAGE: lm_instruct, instruction_type: all)

Epoch 0/299 ━━━╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 277/3600 0:05:06 • 1:01:25 0.90it/s  

Stage2_instruct_t2m (STAGE: lm_instruct, instruction_type: t2m)

Epoch 0/299 ━━━━━━━━━━━╸━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 701/2400 0:12:47 • 0:31:14 0.91it/s 

Stage3 (STAGE: lm_finetune, instruction_type: witht2t)

Epoch 0/299 ━━━━━╺━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 710/5599 0:13:40 • 1:34:16 0.86it/s