Merge pull request #3158 from bghira/perf/dynamo-megacache

Persist and reuse Torch compiler Mega-Caches

Author: bghira

documentation/OPTIONS.es.md

@@ -602,6 +602,23 @@ Para persistir los ajustes en `config.json`, añade las claves equivalentes:
 
 Omite cualquier entrada que quieras heredar de los valores predeterminados de Accelerate (por ejemplo, deja fuera `dynamo_mode` para usar la selección automática).
 
+### `--dynamo_wrapper`
+
+Selecciona el wrapper de host de TorchInductor. `cpp` es el valor predeterminado y reduce el overhead de despacho; `python` conserva el comportamiento de versiones anteriores. La selección forma parte del nombre y manifiesto de Mega-Cache.
+
+### `--dynamo_cache_export`
+
+Ruta opcional para un blob acumulativo de PyTorch Mega-Cache. SimpleTuner carga un blob compatible antes de compilar, lo exporta tras el primer paso exitoso del optimizador y comprueba si hay nuevas claves de artefactos tras cada checkpoint y al finalizar. El manifiesto `<ruta>.manifest.json` registra el entorno PyTorch/Triton/GPU y una suma SHA256. Si una forma no está cubierta, PyTorch compila normalmente y el siguiente exportado incorpora esos artefactos. Usa únicamente cachés de confianza que coincidan con el runtime.
+Si el valor es un directorio, termina en un separador o no tiene extensión, SimpleTuner genera allí un nombre estable basado en el modelo, runtime, acelerador y configuración relevante para el grafo, y busca el mismo nombre en Hub.
+
+### `--dynamo_cache_export_after_first_step`
+
+Cuando es `true` (predeterminado), exporta la Mega-Cache tras el primer paso exitoso. Establece `false` para omitir solo esta exportación temprana; las exportaciones de checkpoints y finalización siguen activas.
+
+### `--dynamo_hub_repo_id`
+
+Repositorio opcional de Hugging Face para recuperar y publicar el blob indicado por `--dynamo_cache_export`. El blob y su manifiesto se publican juntos en un solo commit; un repositorio inexistente se crea como privado. Los fallos de Hub no interrumpen el entrenamiento y la copia local se conserva.
+
 ### `--attention_mechanism`
 
 Se soportan mecanismos de atención alternativos, con distintos niveles de compatibilidad u otros compromisos:

documentation/OPTIONS.hi.md

@@ -601,6 +601,23 @@ Settings को `config.json` में persist करने के लिए 
 
 यदि आप Accelerate defaults inherit करना चाहते हैं तो संबंधित entries छोड़ दें (उदा., `dynamo_mode` न दें ताकि automatic selection उपयोग हो)।
 
+### `--dynamo_wrapper`
+
+TorchInductor का host wrapper चुनता है। `cpp` default है और dispatch overhead घटाता है; `python` पुराने releases का behavior बनाए रखता है। यह चयन Mega-Cache filename और manifest में शामिल होता है।
+
+### `--dynamo_cache_export`
+
+PyTorch Mega-Cache के cumulative blob की optional path। SimpleTuner compilation से पहले compatible blob load करता है, पहले सफल optimizer step के बाद export करता है, और हर checkpoint तथा shutdown पर नई artifact keys की जांच करता है। `<path>.manifest.json` में PyTorch/Triton/GPU runtime और SHA256 दर्ज होते हैं। किसी नई shape के लिए cache entry न मिलने पर PyTorch सामान्य compilation करता है और अगला export नए artifacts जोड़ देता है। केवल trusted cache blobs load करें।
+यदि value directory है, separator पर समाप्त होती है, या उसकी extension नहीं है, तो SimpleTuner model, runtime, accelerator और graph-relevant config से stable filename बनाता है और Hub पर भी वही नाम खोजता है।
+
+### `--dynamo_cache_export_after_first_step`
+
+`true` (default) होने पर पहले सफल optimizer step के बाद Mega-Cache export होती है। केवल इस early export को छोड़ने के लिए `false` करें; checkpoint और final exports सक्रिय रहते हैं।
+
+### `--dynamo_hub_repo_id`
+
+`--dynamo_cache_export` blob को retrieve और publish करने के लिए optional Hugging Face repository। Blob और manifest एक ही commit में upload होते हैं; missing repository private रूप में बनाई जाती है। Hub failure training को abort नहीं करता और local export सुरक्षित रहता है।
+
 ### `--attention_mechanism`
 
 Alternative attention mechanisms समर्थित हैं, जिनके compatibility स्तर या trade‑offs अलग होते हैं:

documentation/OPTIONS.ja.md

@@ -602,6 +602,23 @@ TRAINING_DYNAMO_BACKEND=inductor
 
 Accelerate の既定値を使いたい項目は省略してください（例: 自動選択の `dynamo_mode` を使うなら省略）。
 
+### `--dynamo_wrapper`
+
+TorchInductor の host wrapper を選択します。`cpp` が既定で dispatch overhead を削減し、`python` は以前の release の動作を維持します。この選択は Mega-Cache の filename と manifest に含まれます。
+
+### `--dynamo_cache_export`
+
+累積 PyTorch Mega-Cache blob の任意パスです。SimpleTuner はコンパイル前に互換 blob を読み込み、最初の成功した optimizer step 後に保存し、各 checkpoint と終了時に新しい artifact key を確認します。`<path>.manifest.json` に PyTorch/Triton/GPU runtime と SHA256 を記録します。未対応の shape は通常どおりコンパイルされ、次の export に追加されます。信頼できる cache blob のみ使用してください。
+値が directory、区切り文字で終わる path、または拡張子のない path の場合、model、runtime、accelerator、graph 関連設定から安定した filename を生成し、Hub でも同じ名前を検索します。
+
+### `--dynamo_cache_export_after_first_step`
+
+`true`（既定）の場合、最初の成功した optimizer step 後に Mega-Cache を export します。`false` はこの初回 export のみを無効にし、checkpoint と終了時の export は継続します。
+
+### `--dynamo_hub_repo_id`
+
+`--dynamo_cache_export` blob の取得と公開に使う任意の Hugging Face repository です。Blob と manifest は 1 commit でアップロードされ、存在しない repository は private で作成されます。Hub 障害で学習は中断されず、local export は保持されます。
+
 ### `--attention_mechanism`
 
 代替アテンション機構が利用可能で、互換性やトレードオフが異なります:

documentation/OPTIONS.md

@@ -608,6 +608,41 @@ To persist the settings in `config.json`, add the equivalent keys:
 
 Omit any entries you want to inherit from Accelerate’s defaults (for example, leave out `dynamo_mode` to use automatic selection).
 
+### `--dynamo_wrapper`
+
+Select the host-side wrapper used by TorchInductor. `cpp` is the default and reduces Python dispatch overhead for large compiled regions. Set `python` to retain the wrapper behavior used by earlier SimpleTuner releases. SimpleTuner applies the selection before loading or compiling graphs and includes it in generated Mega-Cache names and manifests.
+
+### `--dynamo_cache_export`
+
+Optional path for a cumulative PyTorch `torch.compile` Mega-Cache blob. At startup, SimpleTuner loads a compatible blob from this path before any model compilation. After the first successful optimizer step it exports the compiled AOTAutograd, Inductor, Triton, autotuning, and PGO artifacts so a later training failure does not discard the cold compile. At normal shutdown it re-exports only when PyTorch reports additional artifact keys, such as those generated for a newly encountered latent shape.
+
+When the value ends in a path separator, identifies an existing directory, or has no filename suffix, SimpleTuner generates a stable filename in that directory. The generated name includes the model family and flavour, accelerator, PyTorch runtime digest, and a digest of graph-relevant configuration such as precision, attention, checkpointing, and LoRA layout. Supplying an explicit filename such as `ltx25-h100.ptcache` uses that name unchanged.
+
+SimpleTuner writes a compatibility manifest beside the blob as `<path>.manifest.json`. The manifest records the exact PyTorch, Triton, CUDA/ROCm, Python, platform, and accelerator identity plus advisory model and compilation configuration signatures. An obvious runtime mismatch is rejected before loading; PyTorch's own cache keys and guards remain authoritative for graph, dtype, stride, and shape compatibility. If no entry covers a batch, training continues with normal runtime compilation and the new artifacts are merged into the next export.
+
+Use runtime-specific paths to avoid replacing a useful cache when changing compiler or hardware versions, for example:
+
+```json
+{
+  "dynamo_backend": "inductor",
+  "dynamo_wrapper": "cpp",
+  "dynamo_use_regional_compilation": true,
+  "dynamo_cache_export": "compiler-caches/ltx25-h100-torch2.11-cu128.ptcache"
+}
+```
+
+Compiler cache blobs contain generated executable code. Load them only from paths and repositories you trust.
+
+### `--dynamo_cache_export_after_first_step`
+
+When `true` (the default), export the Dynamo Mega-Cache immediately after the first successful optimizer step. This protects the expensive initial compile if training later fails. SimpleTuner also checks for new compiler artifact keys after every successful scheduled, manual, rolling, or epoch checkpoint and at normal training completion. Set this option to `false` to skip only the first-step export.
+
+### `--dynamo_hub_repo_id`
+
+Optional Hugging Face model repository used with `--dynamo_cache_export`. SimpleTuner checks the configured relative blob path in this repository before falling back to the local path. If compilation adds artifacts, the blob and manifest are written locally and published together in one Hub commit. A missing repository is created as private. Existing repositories retain their current visibility.
+
+Absolute local export paths use only their filename on the Hub. Relative paths are preserved, so `compiler-caches/ltx25.ptcache` is stored under the same repository subdirectory. For a directory value, SimpleTuner looks for the same generated standard filename locally and on the Hub. Authentication uses `HF_TOKEN`, `HUGGING_FACE_HUB_TOKEN`, or the standard Hugging Face token cache. Hub failures never abort training; the local cache remains available for a later upload attempt.
+
 ### `--attention_mechanism`
 
 Alternative attention mechanisms are supported, with varying levels of compatibility or other trade-offs:

documentation/OPTIONS.pt-BR.md

@@ -600,6 +600,23 @@ Para persistir as configuracoes em `config.json`, adicione as chaves equivalente
 
 Omitir entradas que voce quer herdar dos defaults do Accelerate (por exemplo, deixe `dynamo_mode` ausente para selecao automatica).
 
+### `--dynamo_wrapper`
+
+Seleciona o wrapper host do TorchInductor. `cpp` e o padrao e reduz overhead de dispatch; `python` preserva o comportamento de releases anteriores. A escolha faz parte do filename e manifesto da Mega-Cache.
+
+### `--dynamo_cache_export`
+
+Caminho opcional para um blob cumulativo do PyTorch Mega-Cache. O SimpleTuner carrega um blob compativel antes da compilacao, exporta apos o primeiro passo bem-sucedido do otimizador e verifica novas chaves de artefatos apos cada checkpoint e no encerramento. O manifesto `<caminho>.manifest.json` registra o runtime PyTorch/Triton/GPU e o SHA256. Shapes sem entrada sao compilados normalmente e seus artefatos entram na proxima exportacao. Carregue apenas caches confiaveis.
+Quando o valor e um diretorio, termina em separador ou nao possui extensao, o SimpleTuner gera um filename estavel com base no modelo, runtime, acelerador e configuracao relevante ao grafo, procurando o mesmo nome no Hub.
+
+### `--dynamo_cache_export_after_first_step`
+
+Quando `true` (padrao), exporta a Mega-Cache apos o primeiro passo bem-sucedido. Use `false` para omitir apenas essa exportacao inicial; exportacoes em checkpoints e no final continuam ativas.
+
+### `--dynamo_hub_repo_id`
+
+Repositorio opcional do Hugging Face para recuperar e publicar o blob de `--dynamo_cache_export`. Blob e manifesto sao enviados juntos em um unico commit; um repositorio ausente e criado como privado. Falhas do Hub nao interrompem o treino e a exportacao local e preservada.
+
 ### `--attention_mechanism`
 
 Mecanismos de atencao alternativos sao suportados, com diferentes niveis de compatibilidade e trade-offs:

documentation/OPTIONS.zh.md

@@ -604,6 +604,23 @@ TRAINING_DYNAMO_BACKEND=inductor
 
 省略你希望继承 Accelerate 默认值的项（例如省略 `dynamo_mode` 以使用自动选择）。
 
+### `--dynamo_wrapper`
+
+选择 TorchInductor 的 host wrapper。`cpp` 为默认值并减少调度开销；`python` 保留旧版本行为。该选择会写入 Mega-Cache 文件名和 manifest。
+
+### `--dynamo_cache_export`
+
+累积 PyTorch Mega-Cache blob 的可选路径。SimpleTuner 会在编译前加载兼容 blob，在第一个成功的优化器步骤后导出，并在每次 checkpoint 和关闭时检查新的 artifact key。`<路径>.manifest.json` 会记录 PyTorch/Triton/GPU 运行时和 SHA256。未覆盖的 shape 会正常编译，其产物会加入下一次导出。只应加载可信的缓存 blob。
+当该值是目录、以路径分隔符结尾或没有扩展名时，SimpleTuner 会根据模型、运行时、加速器和影响计算图的配置生成稳定文件名，并在 Hub 上查找同名文件。
+
+### `--dynamo_cache_export_after_first_step`
+
+为 `true`（默认）时，在第一个成功的优化器步骤后导出 Mega-Cache。设为 `false` 只会跳过这次早期导出；checkpoint 和训练结束导出仍会执行。
+
+### `--dynamo_hub_repo_id`
+
+用于获取和发布 `--dynamo_cache_export` blob 的可选 Hugging Face 仓库。Blob 与 manifest 会在一次 commit 中上传；不存在的仓库会创建为 private。Hub 故障不会中止训练，本地导出仍会保留。
+
 ### `--attention_mechanism`
 
 支持多种注意力机制，兼容性与权衡不同：

simpletuner/helpers/models/ltxvideo2/transformer.py

@@ -251,10 +251,17 @@ def _ltx2_prepare_flash_attention_metadata(
         max_seqlen_k = seq_len_kv
     else:
         seqlens_k = keep_mask.sum(dim=1, dtype=torch.int32)
-        if torch.any(seqlens_k == 0):
+        if not torch.compiler.is_compiling() and torch.any(seqlens_k == 0):
             raise ValueError("LTX-2 varlen attention received a mask with no valid key/value tokens.")
-        indices_k = keep_mask.flatten().nonzero(as_tuple=False).flatten()
-        max_seqlen_k = int(seqlens_k.max().item())
+        flat_keep_mask = keep_mask.flatten()
+        valid_destinations = torch.cumsum(flat_keep_mask, dim=0, dtype=torch.int64) - 1
+        invalid_destinations = (
+            flat_keep_mask.sum(dtype=torch.int64) + torch.cumsum(~flat_keep_mask, dim=0, dtype=torch.int64) - 1
+        )
+        destinations = torch.where(flat_keep_mask, valid_destinations, invalid_destinations)
+        source_indices = torch.arange(flat_keep_mask.numel(), device=device)
+        indices_k = torch.empty_like(source_indices).scatter(0, destinations, source_indices)
+        max_seqlen_k = seq_len_kv
 
     cu_seqlens_q = torch.zeros(batch_size + 1, dtype=torch.int32, device=device)
     cu_seqlens_k = torch.zeros(batch_size + 1, dtype=torch.int32, device=device)
@@ -282,6 +289,10 @@ def _ltx2_flash_varlen_hub_attention_prepared(
     else:
         key_packed = key.flatten(0, 1)[metadata.indices_k]
         value_packed = value.flatten(0, 1)[metadata.indices_k]
+        valid_rows = torch.arange(key_packed.shape[0], device=key.device) < metadata.cu_seqlens_k[-1]
+        valid_rows = valid_rows[:, None, None]
+        key_packed = torch.where(valid_rows, key_packed, torch.zeros_like(key_packed))
+        value_packed = torch.where(valid_rows, value_packed, torch.zeros_like(value_packed))
 
     func = _HUB_KERNELS_REGISTRY[AttentionBackendName.FLASH_VARLEN_HUB].kernel_fn
     if func is None:
@@ -305,6 +316,38 @@ def _ltx2_flash_varlen_hub_attention_prepared(
     return output.unflatten(0, (batch_size, seq_len_q))
 
 
+def _ltx2_native_efficient_attention(
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+) -> torch.Tensor:
+    query, key, value = (tensor.permute(0, 2, 1, 3) for tensor in (query, key, value))
+    if attention_mask is not None and attention_mask.dtype == torch.bool:
+        zero = torch.zeros((), dtype=query.dtype, device=query.device)
+        masked = torch.full((), torch.finfo(query.dtype).min, dtype=query.dtype, device=query.device)
+        attention_mask = torch.where(attention_mask, zero, masked)
+    elif attention_mask is not None and attention_mask.dtype != query.dtype:
+        attention_mask = attention_mask.to(query.dtype)
+    if attention_mask is not None and attention_mask.shape[-2] != query.shape[-2]:
+        attention_mask = attention_mask.expand(*attention_mask.shape[:-2], query.shape[-2], attention_mask.shape[-1])
+    if attention_mask is not None and attention_mask.shape[-1] % 8:
+        key_length = attention_mask.shape[-1]
+        padding = 8 - key_length % 8
+        attention_mask = torch.nn.functional.pad(attention_mask, (0, padding))[..., :key_length]
+    output, _, _, _ = torch.ops.aten._scaled_dot_product_efficient_attention(
+        query,
+        key,
+        value,
+        attention_mask,
+        True,
+        dropout_p=0.0,
+        is_causal=False,
+        scale=None,
+    )
+    return output.permute(0, 2, 1, 3)
+
+
 def _ltx2_dispatch_attention(
     query: torch.Tensor,
     key: torch.Tensor,
@@ -333,6 +376,8 @@ def _ltx2_dispatch_attention(
                 raise ValueError("Precomputed LTX-2 Flash Attention metadata does not support context parallelism.")
             return _ltx2_flash_varlen_hub_attention_prepared(query, key, value, flash_attention_metadata)
         attention_mask = _ltx2_normalize_varlen_mask(attention_mask, key.shape[0], key.shape[1])
+    if backend_name == AttentionBackendName._NATIVE_EFFICIENT and parallel_config is None:
+        return _ltx2_native_efficient_attention(query, key, value, attention_mask)
     return dispatch_attention_fn(
         query,
         key,
@@ -2637,7 +2682,8 @@ def prepare_mask(attn, mask, sequence_length, parallel_config=None):
                 for processor in attention_processors
             ) and all(getattr(processor, "_parallel_config", None) is None for processor in attention_processors)
             if use_prepared_flash_attention:
-                _maybe_download_kernel_for_backend(AttentionBackendName.FLASH_VARLEN_HUB)
+                if not torch.compiler.is_compiling():
+                    _maybe_download_kernel_for_backend(AttentionBackendName.FLASH_VARLEN_HUB)
 
                 def prepare_flash_metadata(mask, seq_len_q, seq_len_kv):
                     return _ltx2_prepare_flash_attention_metadata(

simpletuner/helpers/training/checkpointing.py

@@ -13,7 +13,4 @@ def checkpoint(function: Callable[..., Any], *args: Any, **kwargs: Any) -> Any:
     use_reentrant = kwargs.get("use_reentrant")
     if use_reentrant is False and _dynamo_backend_enabled():
         kwargs.setdefault("determinism_check", "none")
-        if torch.compiler.is_compiling():
-            checkpoint_eager = torch.compiler.disable(torch.utils.checkpoint.checkpoint)
-            return checkpoint_eager(function, *args, **kwargs)
     return torch.utils.checkpoint.checkpoint(function, *args, **kwargs)