[ltx-trainer] DDP validation crashes on transformer.num_blocks when STG is enabled

[Kaihua1203]

Summary

ltx-trainer validation crashes under multi-GPU DDP when STG is enabled because ValidationRunner._run_denoising() reads the custom model attribute transformer.num_blocks directly. At validation time, transformer is a torch.nn.parallel.DistributedDataParallel wrapper, which does not expose that attribute.

Training itself runs normally until the first validation interval.

Environment

• LTX-2 commit: 400fd31
• Python: 3.10.13
• PyTorch: 2.10.0+cu128
• Launch: Accelerate multi-GPU, 4 processes
• Hardware: 4 x H100 80 GB
• Training mode: LTX-2.5 LoRA

Minimal configuration

validation:
  interval: 500
  inference_steps: 30
  video_cfg_scale: 3.0
  video_stg_scale: 1.0
  stg_blocks: [28]
  generate_audio: false
  generate_video: true

Launch command:

accelerate launch --multi_gpu --num_processes 4 packages/ltx-trainer/scripts/train.py config.yaml

Error

The run reaches the first validation and then fails in packages/ltx-trainer/src/ltx_trainer/validation_runner.py:

AttributeError: DistributedDataParallel object has no attribute num_blocks

The first failing access is in the STG perturbation setup:

num_blocks=transformer.num_blocks

There is a second equivalent access in the modality-guidance branch:

transformer.num_blocks, device, transformer_dtype

Root cause

Trainer._run_validation() passes self._transformer after Accelerator has wrapped it with DDP. DDP supports forward calls and parameter access, but does not proxy arbitrary custom attributes from the wrapped module. The attribute is available as transformer.module.num_blocks.

Single-GPU validation does not expose the issue because the transformer is not wrapped.

Suggested fix

Resolve the metadata from the inner module while retaining the DDP wrapper for model forward calls:

def _get_transformer_num_blocks(transformer: torch.nn.Module) -> int:
    inner_transformer = getattr(transformer, "module", transformer)
    return int(inner_transformer.num_blocks)

transformer_num_blocks = _get_transformer_num_blocks(transformer)

Then use transformer_num_blocks in both the STG and modality-guidance perturbation builders. X0Model(transformer) should continue receiving the wrapped transformer.

A CPU DDP regression test can reproduce this without GPUs by wrapping a small nn.Module that defines num_blocks and verifying the helper reads it through .module.