torch.compile for Anima Training / Anima 学習における torch.compile

This document explains how to speed up Anima LoRA training (anima_train_network.py) with PyTorch's torch.compile. torch.compile is a just-in-time (JIT) compilation feature that can reduce per-step time by optimizing model execution.

This feature is ported from Musubi Tuner and currently applies to Anima only in sd-scripts. It uses per-block compilation: each Transformer block of the DiT is compiled individually rather than compiling the whole model at once. This keeps the compilation cache small (one artifact reused across the structurally identical blocks), avoids excessive recompilation, and coexists with block swapping.

Note: torch.compile may not work well in every environment. If it fails, simply train without --compile (the traditional path is unchanged).

1. Prerequisites / 前提条件

• Triton is required for torch.compile to be effective. On Windows, install it from the triton-windows project.
• A recent PyTorch (2.x) is required.
• MSVC compiler (Visual Studio 2022 with C++ development tools, or Build Tools 2022) is required on Windows only when using --compile_dynamic true. See Section 6.

2. Command Line Arguments / コマンドライン引数

These arguments are added by anima_train_network.py. They are independent of the legacy --torch_compile option (which routes through accelerate's dynamo and is not the per-block path described here). --compile and --torch_compile cannot be used together.

Basic arguments / 基本的な引数

• --compile: Enable per-block torch.compile for the DiT (requires Triton).
• --compile_backend: Backend to use (default: inductor).
• --compile_mode: Compilation mode (default: default, recommended for training).
  • Choices: default, reduce-overhead, max-autotune, max-autotune-no-cudagraphs
• --compile_dynamic: Dynamic shapes mode (default: None, equivalent to auto). On Windows, true requires the Visual Studio 2022 C++ compiler.
  • Choices: true, false, auto
• --compile_fullgraph: Enable fullgraph mode. Cannot be used with --split_attn.
• --compile_cache_size_limit: Set torch._dynamo.config.cache_size_limit (default: PyTorch default, typically 8-32; recommended: 32).

Additional performance arguments / 追加のパフォーマンス引数

These are independent of --compile and can be used on their own.

• --cuda_allow_tf32: Allow TF32 precision on Ampere or newer GPUs (improves performance).
• --cuda_cudnn_benchmark: Enable cuDNN benchmark mode (may improve performance).

3. Usage Example / 使用例

accelerate launch --num_cpu_threads_per_process 1 --mixed_precision bf16 \
  anima_train_network.py \
  --pretrained_model_name_or_path path/to/anima_dit.safetensors \
  --qwen3 path/to/qwen3 \
  --vae path/to/vae \
  --dataset_config path/to/config.toml \
  --output_dir path/to/output \
  --output_name anima_lora \
  (... other training args ...) \
  --compile \
  --compile_mode default \
  --compile_cache_size_limit 32 \
  --cuda_allow_tf32 \
  --cuda_cudnn_benchmark

On Windows Command Prompt, use ^ instead of \ for line continuation.

4. Notes on Behavior / 挙動に関する注意

• The first iterations are slow. torch.compile performs JIT compilation on the first forward (and backward) pass with a given input shape. This is expected; subsequent steps are much faster.
• Recompilation on shape changes. If --compile_dynamic is not set to true, recompilation occurs whenever an input shape changes. With multi-resolution (aspect-ratio bucketing) datasets, each new bucket resolution triggers a recompilation, so the first epoch can be noticeably slower while later epochs benefit. Increasing --compile_cache_size_limit (e.g. to 32) helps avoid cache eviction across many bucket shapes.
• max-autotune may not work in some cases. If you hit errors, fall back to --compile_mode default.
• --compile_fullgraph may not work depending on the configuration; in particular it cannot be combined with --split_attn (see Section 5).

5. Limitations and Known Issues / 制限事項と既知の問題

• --compile_fullgraph and --split_attn cannot be used together. --split_attn uses dynamic control flow that is incompatible with fullgraph mode. This combination is rejected at startup.
• --compile and --torch_compile cannot be used together. The new per-block path and the legacy accelerate-dynamo path are mutually exclusive and rejected at startup.
• Block swapping (--blocks_to_swap). When block swapping is enabled, the Linear layers inside swapped blocks are automatically excluded from compilation (their weights move between CPU and GPU each step, which conflicts with a compiled graph). The blocks are still compiled, but the speed-up may be smaller than without block swapping.

6. Windows Requirements for --compile_dynamic / Windows での --compile_dynamic の要件

IMPORTANT: On Windows, using --compile_dynamic true requires:

1. Visual Studio 2022 with C++ development tools installed.
2. Either:
   • Running the training script from the "x64 Native Tools Command Prompt for VS 2022", or
   • Running it after executing vcvars64.bat to set the environment variables, e.g. "C:\Program Files (x86)\Microsoft Visual Studio\2022\BuildTools\VC\Auxiliary\Build\vcvars64.bat".

If you hit compilation errors with --compile_dynamic true, make sure you are running from the correct command prompt, or try without --compile_dynamic. See also the PyTorch Inductor Windows documentation.

7. Recommended Settings / 推奨設定

For training, the following is a good starting point:

--compile \
--compile_mode default \
--compile_cache_size_limit 32 \
--cuda_allow_tf32 \
--cuda_cudnn_benchmark

The performance gain depends on the GPU, settings, and dataset. As one example, with batch size 4 and gradient checkpointing enabled, the second-epoch training time improved from about 1:13 to about 1:03, and a 40-step window (steps 60-100) improved from about 24s to about 17s. Block swapping reduces the gain (see Section 5).

--compile \
--compile_mode default \
--compile_cache_size_limit 32 \
--cuda_allow_tf32 \
--cuda_cudnn_benchmark

8. Troubleshooting / トラブルシューティング

• Out of memory. Try a smaller --compile_cache_size_limit, reduce batch size, or use --compile_mode default instead of max-autotune.
• Compilation errors. Ensure Triton is installed correctly. On Windows with --compile_dynamic true, verify the MSVC setup (Section 6). If problems persist, train without --compile.

Additional Resources / 追加リソース

• PyTorch torch.compile tutorial
• PyTorch Inductor Windows documentation
• triton-windows
• Anima LoRA Training Guide