diff --git a/docs/images/timestep_sampling_offset/A1.jpg b/docs/images/timestep_sampling_offset/A1.jpg deleted file mode 100644 index 4a13ce126..000000000 Binary files a/docs/images/timestep_sampling_offset/A1.jpg and /dev/null differ diff --git a/docs/images/timestep_sampling_offset/A2.jpg b/docs/images/timestep_sampling_offset/A2.jpg deleted file mode 100644 index 67039276f..000000000 Binary files a/docs/images/timestep_sampling_offset/A2.jpg and /dev/null differ diff --git a/docs/images/timestep_sampling_offset/A3.jpg b/docs/images/timestep_sampling_offset/A3.jpg deleted file mode 100644 index 3463c2ceb..000000000 Binary files a/docs/images/timestep_sampling_offset/A3.jpg and /dev/null differ diff --git a/docs/images/timestep_sampling_offset/Q1.png b/docs/images/timestep_sampling_offset/Q1.png deleted file mode 100644 index 2851e929d..000000000 Binary files a/docs/images/timestep_sampling_offset/Q1.png and /dev/null differ diff --git a/docs/images/timestep_sampling_offset/Q2.png b/docs/images/timestep_sampling_offset/Q2.png deleted file mode 100644 index ed223ef61..000000000 Binary files a/docs/images/timestep_sampling_offset/Q2.png and /dev/null differ diff --git a/docs/images/timestep_sampling_offset/Q3.png b/docs/images/timestep_sampling_offset/Q3.png deleted file mode 100644 index ed9525f52..000000000 Binary files a/docs/images/timestep_sampling_offset/Q3.png and /dev/null differ diff --git a/docs/images/timestep_sampling_offset/Suppl-Ablation_lora-F1.jpg b/docs/images/timestep_sampling_offset/Suppl-Ablation_lora-F1.jpg deleted file mode 100644 index 20894fd11..000000000 Binary files a/docs/images/timestep_sampling_offset/Suppl-Ablation_lora-F1.jpg and /dev/null differ diff --git a/docs/images/timestep_sampling_offset/Suppl-Ablation_lora-J1.jpg b/docs/images/timestep_sampling_offset/Suppl-Ablation_lora-J1.jpg deleted file mode 100644 index 7d8d3ecd9..000000000 Binary files a/docs/images/timestep_sampling_offset/Suppl-Ablation_lora-J1.jpg and /dev/null differ diff --git a/docs/images/timestep_sampling_offset/Suppl-Ablation_lora-shift_1.jpg b/docs/images/timestep_sampling_offset/Suppl-Ablation_lora-shift_1.jpg deleted file mode 100644 index 3322b51d6..000000000 Binary files a/docs/images/timestep_sampling_offset/Suppl-Ablation_lora-shift_1.jpg and /dev/null differ diff --git a/docs/images/timestep_sampling_offset/histogram.png b/docs/images/timestep_sampling_offset/histogram.png deleted file mode 100644 index 2adf76694..000000000 Binary files a/docs/images/timestep_sampling_offset/histogram.png and /dev/null differ diff --git a/documentation/DATALOADER.es.md b/documentation/DATALOADER.es.md index 8bc0230b4..fdf97a3a9 100644 --- a/documentation/DATALOADER.es.md +++ b/documentation/DATALOADER.es.md @@ -1377,6 +1377,7 @@ Los datasets Webshart cargan shards tar estilo WebDataset mediante el paquete `w - `metadata` es opcional y puede apuntar a metadatos separados con captions. Para repositorios Hugging Face de metadata como `webshart/conceptual-captions-12m-webdataset-metadata`, pasa el repo id; Webshart sigue el layout de subcarpetas del source, como `data/`. - `metadata_backend` debe ser `webshart`; `caption_strategy` debe ser `webshart` o `instanceprompt`. - `webshart.cache_dir` almacena la metadata de SimpleTuner y las caches de Webshart. `shard_cache_gb` y `parallel_downloads` se pasan a la cache de shards de Webshart; define `shard_cache_gb` como `0` para desactivar la cache de shards completos y mantener lecturas por rango indexadas. +- `webshart.caption_key` permite seleccionar campos de caption: usa `"long_caption"` para una clave o `["long_caption", "short_caption"]` para recopilar varias en el orden indicado. Las claves son nombres literales que se buscan en los metadatos JSON de la muestra, después en los metadatos del índice y, finalmente, en las entradas con nombre de sus respectivos diccionarios `captions`; se usa la primera ubicación que contiene la clave. Los valores de texto y listas se convierten en variantes de caption, sin concatenarse en un único prompt. Se ignoran los valores ausentes o vacíos; con `caption_strategy: "webshart"`, se omiten las muestras sin captions seleccionadas aunque tengan captions predeterminadas o un sidecar `.txt`. Omitir la opción conserva la búsqueda predeterminada. Se leen los sidecars JSON cuando su contenido no está en el índice. Las cachés de captions y buckets se separan según el selector configurado. En los ajustes Webshart de la WebUI, introduce una clave por línea. - `webshart_optimize_captions` (grafía alternativa `webshart_optimise_captions`; también se acepta como `optimize_captions`/`optimise_captions` dentro del bloque `webshart`) sondea el layout de captions al arrancar y, cuando los captions residen en miembros tar sidecar `.txt`/`.json` en lugar del índice de metadata, los consolida una sola vez en la cache local de metadata de Webshart. Sin esta opción, los datasets con captions en sidecars (por ejemplo `laion/conceptual-captions-12m-webdataset`) pagan una lectura por rango por muestra cada vez que se enumeran los captions — al arrancar, al guardar checkpoints y al generar la model card. Los datasets cuya metadata ya incluye los captions omiten la consolidación automáticamente. #### Optimizar captions por adelantado diff --git a/documentation/DATALOADER.hi.md b/documentation/DATALOADER.hi.md index 999910cc9..5db48b099 100644 --- a/documentation/DATALOADER.hi.md +++ b/documentation/DATALOADER.hi.md @@ -1377,6 +1377,7 @@ Webshart datasets `webshart` package के जरिए WebDataset-style tar sh - `metadata` optional है और captions वाले separate metadata location को point कर सकता है। `webshart/conceptual-captions-12m-webdataset-metadata` जैसे Hugging Face metadata repos के लिए repo id दें; Webshart source shard के `data/` जैसे subfolder layout को follow करता है। - `metadata_backend` को `webshart` होना चाहिए; `caption_strategy` `webshart` या `instanceprompt` हो सकता है। - `webshart.cache_dir` SimpleTuner metadata और Webshart caches store करता है। `shard_cache_gb` और `parallel_downloads` Webshart shard cache को pass किए जाते हैं; whole-shard caching disable करने और indexed range reads बनाए रखने के लिए `shard_cache_gb` को `0` सेट करें। +- `webshart.caption_key` से caption फ़ील्ड चुन सकते हैं: एक कुंजी के लिए `"long_caption"` या कई कुंजियों को क्रम से लेने के लिए `["long_caption", "short_caption"]` दें। कुंजियाँ सीधे नाम के रूप में मिलाई जाती हैं: पहले sample के JSON metadata में, फिर indexed metadata में, और अंत में उनके `captions` dictionaries की नामित entries में। कुंजी जिस पहले स्थान पर मिलती है, वही उपयोग होता है। String और list के मान अलग caption विकल्प बनते हैं, एक prompt में जोड़े नहीं जाते। अनुपस्थित या खाली मान छोड़ दिए जाते हैं; `caption_strategy: "webshart"` के साथ चुने हुए captions न होने पर sample छोड़ दिया जाता है, भले ही उसमें default caption या `.txt` sidecar हो। यह विकल्प न देने पर default lookup बना रहता है। JSON सामग्री index में न होने पर JSON sidecar पढ़ा जाता है। Caption और bucket caches चुनी गई कुंजियों के अनुसार अलग रखे जाते हैं। WebUI की Webshart settings में हर पंक्ति पर एक कुंजी लिखें। - `webshart_optimize_captions` (alternate spelling `webshart_optimise_captions`; `webshart` block के अंदर `optimize_captions`/`optimise_captions` भी accepted हैं) startup पर caption layout probe करता है और, जब captions metadata index की बजाय `.txt`/`.json` sidecar tar members में हों, उन्हें एक बार local Webshart metadata cache में fold कर देता है। इसके बिना, sidecar-caption datasets (जैसे `laion/conceptual-captions-12m-webdataset`) को हर बार captions enumerate होने पर — startup, checkpointing और model card generation में — प्रति sample एक range read की कीमत चुकानी पड़ती है। जिन datasets की metadata में captions पहले से embedded हैं, वे coalescing अपने आप skip कर देते हैं। #### Captions को पहले से optimize करना diff --git a/documentation/DATALOADER.ja.md b/documentation/DATALOADER.ja.md index 5753389d1..a304ab320 100644 --- a/documentation/DATALOADER.ja.md +++ b/documentation/DATALOADER.ja.md @@ -1378,6 +1378,7 @@ Webshart データセットは `webshart` パッケージで WebDataset 形式 - `metadata` は任意で、captions を含む別 metadata location を指定できます。`webshart/conceptual-captions-12m-webdataset-metadata` のような Hugging Face metadata repo では repo id だけを渡します。Webshart は source shard の `data/` などのサブフォルダ構成に従います。 - `metadata_backend` は `webshart`、`caption_strategy` は `webshart` または `instanceprompt` にします。 - `webshart.cache_dir` は SimpleTuner metadata と Webshart caches を保存します。`shard_cache_gb` と `parallel_downloads` は Webshart の shard cache に渡されます。`shard_cache_gb` を `0` にすると、shard 全体の cache を無効にし、index 付き range read を維持します。 +- `webshart.caption_key` でキャプションのフィールドを選択できます。1つなら `"long_caption"`、複数なら `["long_caption", "short_caption"]` を指定すると、その順で収集します。キーはパスではなくそのままの名前として、サンプルの JSON メタデータ、インデックスのメタデータ、それぞれの `captions` 辞書内の名前付き項目の順に検索し、最初にキーが見つかった場所を使います。文字列やリストの値は1つのプロンプトに連結せず、キャプション候補になります。欠落した値や空の値は無視されます。`caption_strategy: "webshart"` では、選択したキャプションがないサンプルは、既定のキャプションや `.txt` サイドカーがあってもスキップされます。省略すると既定の取得方法を維持します。JSON の内容がインデックスにない場合はサイドカーを読み込みます。キャプションとバケットのキャッシュは設定したキーに応じて分離されます。WebUI の Webshart 設定では1行に1つのキーを入力します。 - `webshart_optimize_captions`（別綴り `webshart_optimise_captions`。`webshart` ブロック内では `optimize_captions`/`optimise_captions` も受け付けます）は起動時に caption layout を probe し、captions が metadata index ではなく `.txt`/`.json` の sidecar tar member にある場合、それらをローカルの Webshart metadata cache に一度だけ統合します。このオプションがないと、sidecar caption の dataset（たとえば `laion/conceptual-captions-12m-webdataset`）は captions を列挙するたび — 起動時、checkpoint 時、model card 生成時 — にサンプルごとに 1 回の range read が発生します。metadata に captions が既に埋め込まれている dataset では、統合は自動的にスキップされます。 #### captions を事前に最適化する diff --git a/documentation/DATALOADER.md b/documentation/DATALOADER.md index 87f43e5d4..fe7f51e73 100644 --- a/documentation/DATALOADER.md +++ b/documentation/DATALOADER.md @@ -1433,6 +1433,7 @@ Webshart datasets load WebDataset-style tar shards through the `webshart` packag - `metadata_backend` must be `webshart`; it reads dimensions and captions from Webshart metadata. - `caption_strategy` should be `webshart` to train from metadata captions, or `instanceprompt` to ignore stored captions. - `webshart.cache_dir` stores SimpleTuner metadata plus Webshart metadata and shard caches. `shard_cache_gb` and `parallel_downloads` are passed to Webshart's shard cache; set `shard_cache_gb` to `0` to disable whole-shard caching and retain indexed range reads. +- `webshart.caption_key` optionally selects caption fields: use `"long_caption"` for one key or `["long_caption", "short_caption"]` to collect multiple keys in order. Keys are literal names, checked in the sample’s JSON metadata, then its indexed metadata, then named entries inside their `captions` dictionaries; the first location containing a key wins. String and list values become caption variants, rather than being joined into one prompt. Missing or empty values are ignored; with `caption_strategy: "webshart"`, samples with no selected captions are skipped, even if they have default captions or a `.txt` sidecar. Omitting the option keeps the default caption lookup. JSON sidecars are read when their contents are absent from the index. Caption and bucket caches are separated by the configured selector. In the WebUI’s Webshart settings, enter one key per line. - `webshart_optimize_captions` (alt spelling `webshart_optimise_captions`; also accepted as `optimize_captions`/`optimise_captions` inside the `webshart` block) probes the caption layout at startup and, when captions live in `.txt`/`.json` sidecar tar members rather than the metadata index, folds them into the local Webshart metadata cache once. Without it, sidecar-caption datasets (for example `laion/conceptual-captions-12m-webdataset`) pay one range read per sample every time captions are enumerated — startup, checkpointing, and model card generation. Datasets whose metadata already embeds captions skip the coalescing automatically. #### Optimizing captions ahead of time diff --git a/documentation/DATALOADER.pt-BR.md b/documentation/DATALOADER.pt-BR.md index f9a77722a..06029b9cd 100644 --- a/documentation/DATALOADER.pt-BR.md +++ b/documentation/DATALOADER.pt-BR.md @@ -1377,6 +1377,7 @@ Datasets Webshart carregam shards tar no estilo WebDataset pelo pacote `webshart - `metadata` é opcional e pode apontar para metadados separados com captions. Para repos Hugging Face de metadata como `webshart/conceptual-captions-12m-webdataset-metadata`, passe o repo id; o Webshart segue o layout de subpastas do source, como `data/`. - `metadata_backend` deve ser `webshart`; `caption_strategy` deve ser `webshart` ou `instanceprompt`. - `webshart.cache_dir` armazena os metadados do SimpleTuner e os caches do Webshart. `shard_cache_gb` e `parallel_downloads` são passados ao cache de shards do Webshart; defina `shard_cache_gb` como `0` para desativar o cache de shards completos e manter leituras por intervalo indexadas. +- `webshart.caption_key` permite selecionar campos de caption: use `"long_caption"` para uma chave ou `["long_caption", "short_caption"]` para coletar várias na ordem indicada. As chaves são nomes literais, buscados nos metadados JSON da amostra, depois nos metadados do índice e, por fim, nas entradas nomeadas dos respectivos dicionários `captions`; vale o primeiro local que contém a chave. Valores de texto e listas tornam-se variantes de caption, sem serem concatenados em um único prompt. Valores ausentes ou vazios são ignorados; com `caption_strategy: "webshart"`, amostras sem captions selecionadas são ignoradas mesmo que tenham captions padrão ou um sidecar `.txt`. Omitir a opção mantém a busca padrão. Sidecars JSON são lidos quando seu conteúdo não está no índice. As caches de captions e buckets são separadas conforme o seletor configurado. Nas configurações Webshart da WebUI, insira uma chave por linha. - `webshart_optimize_captions` (grafia alternativa `webshart_optimise_captions`; também aceito como `optimize_captions`/`optimise_captions` dentro do bloco `webshart`) sonda o layout de captions na inicialização e, quando os captions ficam em membros tar sidecar `.txt`/`.json` em vez do índice de metadados, consolida-os uma única vez no cache local de metadados do Webshart. Sem essa opção, datasets com captions em sidecars (por exemplo `laion/conceptual-captions-12m-webdataset`) pagam uma leitura por intervalo por amostra sempre que os captions são enumerados — na inicialização, nos checkpoints e na geração do model card. Datasets cujos metadados já embutem os captions pulam a consolidação automaticamente. #### Otimizando captions com antecedência diff --git a/documentation/DATALOADER.zh.md b/documentation/DATALOADER.zh.md index f1d3d53f7..cfddca5e7 100644 --- a/documentation/DATALOADER.zh.md +++ b/documentation/DATALOADER.zh.md @@ -1377,6 +1377,7 @@ Webshart 数据集通过 `webshart` 包加载 WebDataset 风格的 tar shards。 - `metadata` 可选，可指向包含 captions 的独立 metadata location。对于 `webshart/conceptual-captions-12m-webdataset-metadata` 这样的 Hugging Face metadata repo，传 repo id 即可；Webshart 会跟随 source shard 的 `data/` 等子目录布局。 - `metadata_backend` 必须为 `webshart`；`caption_strategy` 应为 `webshart` 或 `instanceprompt`。 - `webshart.cache_dir` 存储 SimpleTuner metadata 与 Webshart caches。`shard_cache_gb` 和 `parallel_downloads` 会传给 Webshart 的 shard cache；将 `shard_cache_gb` 设为 `0` 可禁用整 shard cache，并保留基于索引的 range reads。 +- `webshart.caption_key` 可用于选择字幕字段：单个键使用 `"long_caption"`，多个键使用 `["long_caption", "short_caption"]`，按列表顺序收集。键按字面名称匹配，依次检查样本的 JSON 元数据、索引元数据及两者 `captions` 字典中的命名条目；采用第一个包含该键的位置。字符串和列表值作为字幕候选，不会拼接成一个提示词。缺失或空值会被忽略；使用 `caption_strategy: "webshart"` 时，没有选中字幕的样本将被跳过，即使它有默认字幕或 `.txt` 伴随文件。省略此选项将保留默认查找方式。如果索引未包含 JSON 内容，则读取 JSON 伴随文件。字幕和分桶缓存按配置的键分别保存。在 WebUI 的 Webshart 设置中，每行输入一个键。 - `webshart_optimize_captions`（另一拼写 `webshart_optimise_captions`；在 `webshart` 块内也接受 `optimize_captions`/`optimise_captions`）会在启动时探测 caption 布局，当 captions 位于 `.txt`/`.json` sidecar tar 成员中而不是 metadata 索引中时，将它们一次性合并进本地 Webshart metadata cache。若不启用，sidecar caption 数据集（例如 `laion/conceptual-captions-12m-webdataset`）在每次枚举 captions 时——启动、checkpoint、生成 model card——都要为每个样本付出一次 range read。metadata 中已内嵌 captions 的数据集会自动跳过合并。 #### 提前优化 captions diff --git a/documentation/OPTIONS.es.md b/documentation/OPTIONS.es.md index e003bb31c..bfa3316c1 100644 --- a/documentation/OPTIONS.es.md +++ b/documentation/OPTIONS.es.md @@ -875,6 +875,7 @@ Muchas configuraciones se establecen a través del [dataloader config](DATALOADE ### `--post_upload_script` +- Tras cada validación integrada completada (manual desde la WebUI, programada o de referencia del modelo base), el script también se ejecuta si no hay un proveedor de publicación configurado. En este caso, `{local_checkpoint_path}` es `output_dir` (no se requiere un checkpoint) y `{remote_checkpoint_path}` está vacío. Los proveedores configurados mantienen un hook por subida exitosa, con las rutas devueltas; las subidas fallidas no activan un hook local. Las validaciones omitidas, fallidas o canceladas y los lanzamientos de scripts externos no activan este hook de finalización. - **Qué**: Ejecutable opcional que se ejecuta después de que cada proveedor de publicación y la subida a Hugging Face Hub termina (subidas finales del modelo y de checkpoints). Se ejecuta de forma asíncrona para que el entrenamiento no se bloquee. - **Marcadores**: Mismas sustituciones que `--validation_external_script`, además de `{remote_checkpoint_path}` (URI devuelta por el proveedor) para que puedas reenviar la URL publicada a sistemas downstream. - **Notas**: @@ -1989,10 +1990,11 @@ Mapeo de opciones upstream (LayerSync → SimpleTuner): ### `--disk_low_threshold` -- **Qué**: Espacio mínimo libre en disco requerido antes de guardar checkpoints. -- **Por qué**: Previene que el entrenamiento falle por errores de disco lleno al detectar espacio bajo tempranamente y tomar una acción configurada. +- **Qué**: Espacio mínimo libre en disco requerido antes de guardar checkpoints y compilar grafos con TorchInductor (incluida la recompilación). +- **Por qué**: Detecta poco espacio anticipadamente y aplica la acción configurada al sistema de archivos de los checkpoints o de la caché del compilador. - **Formato**: Cadena de tamaño como `100G`, `50M`, `1T`, `500K`, o bytes simples. - **Por defecto**: Ninguno (función desactivada) +- **Alcance**: Cada proceso que compila comprueba `TORCHINDUCTOR_CACHE_DIR` (o el valor predeterminado de PyTorch) y un `TRITON_CACHE_DIR` separado si está configurado. Es una comprobación previa, no una reserva: la compilación puede agotar el disco tras iniciarse y la inicialización del backend puede escribir antes de la comprobación. Elija un umbral suficiente para compilar. Estas comprobaciones no recuperan errores de memoria CUDA ni reintentan compilaciones fallidas. ### `--disk_low_action` diff --git a/documentation/OPTIONS.hi.md b/documentation/OPTIONS.hi.md index a87382f05..75e4273cf 100644 --- a/documentation/OPTIONS.hi.md +++ b/documentation/OPTIONS.hi.md @@ -874,6 +874,7 @@ Alternative attention mechanisms समर्थित हैं, जिनक ### `--post_upload_script` +- हर पूर्ण अंतर्निहित सत्यापन (WebUI से मैन्युअल, निर्धारित या बेस मॉडल बेंचमार्क) के बाद, कोई प्रकाशन प्रदाता कॉन्फ़िगर न होने पर भी स्क्रिप्ट चलती है। इस स्थिति में `{local_checkpoint_path}` का मान `output_dir` होता है (चेकपॉइंट आवश्यक नहीं है), और `{remote_checkpoint_path}` खाली होता है। कॉन्फ़िगर किए गए प्रदाताओं के लिए हर सफल अपलोड पर लौटाए गए पथों के साथ एक हुक चलता है; असफल अपलोड स्थानीय हुक नहीं चलाते। छोड़े गए, असफल या रद्द किए गए सत्यापन और बाहरी स्क्रिप्ट का लॉन्च इस पूर्णता हुक को नहीं चलाते। - **What**: हर publishing provider और Hugging Face Hub upload (final model और checkpoint uploads) के बाद optional executable चलता है। यह asynchronous चलता है ताकि training block न हो। - **Placeholders**: `--validation_external_script` जैसे replacements, साथ ही `{remote_checkpoint_path}` (provider द्वारा लौटाया गया URI) ताकि आप published URL downstream systems को भेज सकें। - **Notes**: @@ -1986,10 +1987,11 @@ Upstream option mapping (LayerSync → SimpleTuner): ### `--disk_low_threshold` -- **What**: checkpoint saves से पहले आवश्यक न्यूनतम खाली disk space। -- **Why**: disk full errors से training crash को रोकता है, कम space का जल्दी पता लगाकर configured action लेता है। +- **What**: checkpoint saves और TorchInductor graph compilation (दोबारा compilation सहित) से पहले आवश्यक न्यूनतम खाली disk space। +- **Why**: कम space का जल्दी पता लगाकर checkpoint या compiler cache के filesystem पर configured action लागू करता है। - **Format**: size string जैसे `100G`, `50M`, `1T`, `500K`, या plain bytes। - **Default**: None (feature disabled) +- **Scope**: compilation करने वाला हर process `TORCHINDUCTOR_CACHE_DIR` (या PyTorch की default directory) और configured होने पर अलग `TRITON_CACHE_DIR` की जाँच करता है। यह केवल पहले की जाँच है, space आरक्षित नहीं करता: compilation शुरू होने के बाद भी disk भर सकती है और backend initialization जाँच से पहले लिख सकता है। compilation के लिए पर्याप्त threshold चुनें। यह CUDA out-of-memory errors से recovery या failed compilation को retry नहीं करता। ### `--disk_low_action` diff --git a/documentation/OPTIONS.ja.md b/documentation/OPTIONS.ja.md index d423f490c..2529bf314 100644 --- a/documentation/OPTIONS.ja.md +++ b/documentation/OPTIONS.ja.md @@ -876,6 +876,7 @@ TorchInductor の host wrapper を選択します。`cpp` が既定で dispatch ### `--post_upload_script` +- 組み込み検証が完了するたびに（WebUI からの手動実行、定期実行、ベースモデルのベンチマーク）、公開プロバイダーが未設定の場合もスクリプトが実行されます。この場合、`{local_checkpoint_path}` は `output_dir`（チェックポイントは不要）、`{remote_checkpoint_path}` は空になります。設定済みのプロバイダーでは、成功したアップロードごとに返されたパスでフックが実行され、アップロード失敗時にローカルフックは実行されません。スキップ・失敗・中止された検証や外部スクリプトの起動では、この完了フックは実行されません。 - **内容**: 各公開先および Hugging Face Hub へのアップロード完了後に実行される任意の実行ファイル（最終モデルとチェックポイントのアップロード）。学習をブロックしないよう非同期で実行されます。 - **プレースホルダ**: `--validation_external_script` と同じ置換に加えて `{remote_checkpoint_path}`（プロバイダが返す URI）が使えます。公開 URL を下流システムへ渡す用途に便利です。 - **注記**: @@ -1989,10 +1990,11 @@ LayerSync は同一 Transformer 内の「学生」レイヤーを、より強い ### `--disk_low_threshold` -- **内容**: チェックポイント保存前に必要な最小空きディスク容量。 -- **理由**: ディスク容量不足を早期に検知して設定されたアクションを実行することで、ディスク満杯エラーによる学習クラッシュを防止します。 +- **内容**: チェックポイント保存および TorchInductor グラフコンパイル（再コンパイルを含む）の前に必要な最小空きディスク容量。 +- **理由**: 容量不足を早期に検知し、チェックポイントまたはコンパイラキャッシュのファイルシステムに対して設定されたアクションを実行します。 - **形式**: `100G`、`50M`、`1T`、`500K` のようなサイズ文字列、またはバイト数。 - **デフォルト**: なし（機能無効） +- **対象**: コンパイルする各プロセスで `TORCHINDUCTOR_CACHE_DIR`（未設定なら PyTorch の既定値）と、設定されている場合は別の `TRITON_CACHE_DIR` を確認します。これは事前確認であり、容量の予約ではありません。コンパイル開始後に容量が不足する場合や、確認前にバックエンドの初期化が書き込む場合があります。コンパイルに十分な閾値を設定してください。CUDA メモリ不足からの復旧や失敗したコンパイルの再試行は行いません。 ### `--disk_low_action` diff --git a/documentation/OPTIONS.md b/documentation/OPTIONS.md index 7d4c2826d..a93e11d23 100644 --- a/documentation/OPTIONS.md +++ b/documentation/OPTIONS.md @@ -899,6 +899,7 @@ A lot of settings are instead set through the [dataloader config](DATALOADER.md) ### `--post_upload_script` +- After each completed built-in validation (manual WebUI, scheduled, or base-model benchmark), the script also runs when no publishing provider is configured. In this case, `{local_checkpoint_path}` is `output_dir` (no checkpoint is required), and `{remote_checkpoint_path}` is empty. Configured providers retain one hook per successful upload, with their returned paths; failed uploads do not trigger a local hook. Skipped, failed, or aborted validation and external-script launches do not trigger this completion hook. - **What**: Optional executable run after each publishing provider and Hugging Face Hub upload finishes (final model and checkpoint uploads). Runs asynchronously so training doesn't block. - **Placeholders**: Same replacements as `--validation_external_script`, plus `{remote_checkpoint_path}` (URI returned by the provider) so you can forward the published URL to downstream systems. - **Notes**: @@ -2011,10 +2012,11 @@ Upstream option mapping (LayerSync → SimpleTuner): ### `--disk_low_threshold` -- **What**: Minimum free disk space required before checkpoint saves. -- **Why**: Prevents training from crashing due to disk full errors during checkpoint saves by detecting low space early and taking a configured action. +- **What**: Minimum free disk space required before checkpoint saves and TorchInductor graph compilation (including recompilation). +- **Why**: Detects low space early and applies the configured action to the checkpoint or compiler cache filesystem. - **Format**: Size string like `100G`, `50M`, `1T`, `500K`, or plain bytes. - **Default**: None (feature disabled) +- **Scope**: Compilation checks use `TORCHINDUCTOR_CACHE_DIR` (or PyTorch’s default) and a separate `TRITON_CACHE_DIR` when configured, on every compiling process. This is a preflight check, not a space reservation: compilation can still exhaust disk space after it starts, and backend initialization may write before the check. Choose a threshold large enough for compilation. Disk checks do not recover CUDA out-of-memory errors or retry failed compilation. ### `--disk_low_action` diff --git a/documentation/OPTIONS.pt-BR.md b/documentation/OPTIONS.pt-BR.md index 2acce8d0f..070d18b61 100644 --- a/documentation/OPTIONS.pt-BR.md +++ b/documentation/OPTIONS.pt-BR.md @@ -873,6 +873,7 @@ Muitas configuracoes sao definidas no [dataloader config](DATALOADER.md), mas es ### `--post_upload_script` +- Após cada validação integrada concluída (manual pela WebUI, agendada ou benchmark do modelo base), o script também é executado quando nenhum provedor de publicação está configurado. Nesse caso, `{local_checkpoint_path}` é `output_dir` (não é necessário um checkpoint), e `{remote_checkpoint_path}` fica vazio. Os provedores configurados mantêm um hook por upload bem-sucedido, com os caminhos retornados; uploads com falha não acionam um hook local. Validações ignoradas, com falha ou canceladas e a inicialização de scripts externos não acionam esse hook de conclusão. - **O que**: Executavel opcional rodado apos cada provedor de publicacao e upload no Hugging Face Hub terminar (modelo final e uploads de checkpoints). Roda de forma assincrona para nao bloquear o treinamento. - **Placeholders**: Mesmas substituicoes de `--validation_external_script`, mais `{remote_checkpoint_path}` (URI retornada pelo provedor) para que voce encaminhe a URL publicada para sistemas downstream. - **Notas**: @@ -1984,10 +1985,11 @@ Mapeamento de opcoes upstream (LayerSync → SimpleTuner): ### `--disk_low_threshold` -- **O que**: Espaco minimo livre em disco necessario antes de salvar checkpoints. -- **Por que**: Previne falhas no treinamento por erros de disco cheio detectando espaco baixo antecipadamente e tomando uma acao configurada. +- **O que**: Espaço mínimo livre em disco necessário antes de salvar checkpoints e compilar grafos com TorchInductor (incluindo recompilação). +- **Por que**: Detecta pouco espaço antecipadamente e aplica a ação configurada ao sistema de arquivos dos checkpoints ou do cache do compilador. - **Formato**: String de tamanho como `100G`, `50M`, `1T`, `500K`, ou bytes simples. - **Padrao**: Nenhum (funcionalidade desativada) +- **Escopo**: Cada processo que compila verifica `TORCHINDUCTOR_CACHE_DIR` (ou o padrão do PyTorch) e um `TRITON_CACHE_DIR` separado, quando configurado. É uma verificação prévia, não uma reserva: a compilação ainda pode esgotar o disco após iniciar, e a inicialização do backend pode escrever antes da verificação. Escolha um limite suficiente para compilar. As verificações não recuperam erros de falta de memória CUDA nem repetem compilações que falharam. ### `--disk_low_action` diff --git a/documentation/OPTIONS.zh.md b/documentation/OPTIONS.zh.md index 79b330d0d..0428ae3c0 100644 --- a/documentation/OPTIONS.zh.md +++ b/documentation/OPTIONS.zh.md @@ -878,6 +878,7 @@ TRAINING_DYNAMO_BACKEND=inductor ### `--post_upload_script` +- 每次内置验证完成后（WebUI 手动验证、定时验证或基础模型基准测试），如果未配置发布提供者，也会运行此脚本。此时，`{local_checkpoint_path}` 为 `output_dir`（无需检查点），`{remote_checkpoint_path}` 为空。已配置的提供者仍在每次成功上传后运行一次钩子，并使用其返回的路径；上传失败不会触发本地钩子。跳过、失败或中止的验证，以及外部脚本的启动，都不会触发此完成钩子。 - **内容**：每个发布目标与 Hugging Face Hub 上传完成后运行的可执行脚本（最终模型与检查点上传）。异步运行，不阻塞训练。 - **占位符**：与 `--validation_external_script` 相同的替换，加上 `{remote_checkpoint_path}`（提供者返回的 URI），可用于将发布 URL 传递给下游系统。 - **说明**： @@ -1991,10 +1992,11 @@ LayerSync 通过在同一 Transformer 内让“学生”层对齐更强的“教 ### `--disk_low_threshold` -- **内容**：检查点保存前所需的最小可用磁盘空间。 -- **原因**：通过提前检测磁盘空间不足并采取配置的操作，防止训练因磁盘已满错误而崩溃。 +- **内容**：检查点保存和 TorchInductor 图编译（包括重新编译）前所需的最小可用磁盘空间。 +- **原因**：提前检测空间不足，并对检查点或编译器缓存所在的文件系统执行配置的操作。 - **格式**：大小字符串，如 `100G`、`50M`、`1T`、`500K`，或纯字节数。 - **默认**：无（功能禁用） +- **范围**：每个执行编译的进程都会检查 `TORCHINDUCTOR_CACHE_DIR`（或 PyTorch 默认目录），以及配置的独立 `TRITON_CACHE_DIR`。这只是预检查，不会预留空间：编译开始后仍可能耗尽空间，后端初始化也可能在检查前写入。请设置足够支持编译的阈值。磁盘检查不会恢复 CUDA 内存不足错误，也不会重试失败的编译。 ### `--disk_low_action` diff --git a/setup.py b/setup.py index f0808b0c7..5e3dd6bc2 100644 --- a/setup.py +++ b/setup.py @@ -351,7 +351,7 @@ def _collect_package_files(*directories: str): "aiosqlite>=0.19.0", "httpx>=0.28.0", "psutil>=5.9.0", - "webshart>=0.5.3", + "webshart>=0.5.4", ] # Nightly CUDA extras contain direct URLs that PyPI rejects, so only include them diff --git a/simpletuner/__init__.py b/simpletuner/__init__.py index a9fe2696d..1cf441c40 100644 --- a/simpletuner/__init__.py +++ b/simpletuner/__init__.py @@ -126,4 +126,4 @@ def _suppress_swigvarlink(message, *args, **kwargs): warnings.warn = _suppress_swigvarlink -__version__ = "4.9.1" +__version__ = "4.9.2" diff --git a/simpletuner/helpers/data_backend/builders/webshart.py b/simpletuner/helpers/data_backend/builders/webshart.py index 8795606aa..020950893 100644 --- a/simpletuner/helpers/data_backend/builders/webshart.py +++ b/simpletuner/helpers/data_backend/builders/webshart.py @@ -55,6 +55,7 @@ def _create_backend(self, config: BaseBackendConfig) -> WebshartDataBackend: "compress_cache": self._get_compression_setting(config), "dataset_type": getattr(config, "dataset_type", "image"), "optimize_captions": bool(getattr(config, "webshart_optimize_captions", None) or False), + "caption_key": config.webshart_caption_key, } if is_mock_backend: backend_kwargs["identifier"] = config.id diff --git a/simpletuner/helpers/data_backend/config/image.py b/simpletuner/helpers/data_backend/config/image.py index 45d2e98a9..3bf17b2c3 100644 --- a/simpletuner/helpers/data_backend/config/image.py +++ b/simpletuner/helpers/data_backend/config/image.py @@ -63,6 +63,7 @@ class ImageBackendConfig(BaseBackendConfig): webshart_buffer_size: Optional[int] = None webshart_max_file_size: Optional[int] = None webshart_optimize_captions: Optional[bool] = None + webshart_caption_key: Optional[Union[str, List[str]]] = None vae_cache_clear_each_epoch: Optional[bool] = None probability: float = 1.0 @@ -187,6 +188,8 @@ def _get_arg(key: str, default: Any = None) -> Any: if config.backend_type == "webshart": webshart_block = backend_dict.get("webshart", {}) or {} config.webshart = webshart_block + config.webshart_caption_key = webshart_block.get("caption_key") + validators.validate_webshart_caption_key(config.webshart_caption_key) config.webshart_source = backend_dict.get("source", webshart_block.get("source")) config.webshart_metadata = backend_dict.get("metadata", webshart_block.get("metadata")) config.webshart_hf_token = backend_dict.get("hf_token", webshart_block.get("hf_token")) @@ -570,6 +573,8 @@ def to_dict(self) -> Dict[str, Any]: webshart_config["max_file_size"] = self.webshart_max_file_size if self.webshart_optimize_captions is not None: webshart_config["optimize_captions"] = self.webshart_optimize_captions + if self.webshart_caption_key is not None: + webshart_config["caption_key"] = self.webshart_caption_key if self.video is not None: config["video"] = self.video diff --git a/simpletuner/helpers/data_backend/config/validators.py b/simpletuner/helpers/data_backend/config/validators.py index 8fa9b8a0a..1f1f1a168 100644 --- a/simpletuner/helpers/data_backend/config/validators.py +++ b/simpletuner/helpers/data_backend/config/validators.py @@ -145,6 +145,14 @@ def validate_huggingface_backend_settings( return {"metadata_backend": metadata_backend or "huggingface", "caption_strategy": caption_strategy or "huggingface"} +def validate_webshart_caption_key(caption_key) -> None: + if caption_key is None: + return + keys = [caption_key] if isinstance(caption_key, str) else caption_key + if not isinstance(keys, list) or not keys or any(not isinstance(key, str) or not key.strip() for key in keys): + raise ValueError("webshart.caption_key must be a non-empty string or a non-empty list of non-empty strings.") + + def validate_webshart_backend_settings( backend_type: str, metadata_backend: Optional[str], caption_strategy: Optional[str], backend_id: str ) -> Dict[str, str]: diff --git a/simpletuner/helpers/data_backend/factory.py b/simpletuner/helpers/data_backend/factory.py index 606eb9c33..8632d7291 100644 --- a/simpletuner/helpers/data_backend/factory.py +++ b/simpletuner/helpers/data_backend/factory.py @@ -10,7 +10,7 @@ from math import sqrt from pathlib import Path from types import SimpleNamespace -from typing import Any, Dict, Iterable, List, Optional, Tuple, Union +from typing import Any, Dict, List, Optional, Tuple, Union from .runtime import BatchFetcher @@ -66,26 +66,6 @@ def _normalise_vae_cache_config(config: Dict[str, Any]) -> Tuple[bool, bool]: return vae_cache_disable, vae_cache_ondemand -def _coerce_bucket_keys(indices: Dict[Any, Iterable]) -> Dict[Any, list]: - """Return a copy of aspect ratio bucket indices with numeric keys coerced to float.""" - coerced: Dict[Any, list] = {} - for key, values in (indices or {}).items(): - try: - coerced_key: Any = float(key) - except (TypeError, ValueError): - coerced_key = key - if isinstance(values, dict): - iterable_values = [values] - elif isinstance(values, str): - iterable_values = [values] - elif isinstance(values, Iterable): - iterable_values = list(values) - else: - iterable_values = [values] - coerced.setdefault(coerced_key, []).extend(iterable_values) - return coerced - - import numpy as np import pandas as pd import torch @@ -3179,8 +3159,6 @@ def _configure_metadata_backend(self, backend: Dict[str, Any], init_backend: Dic # Restore the live-authoritative runtime config after metadata cache loading. StateTracker.set_data_backend_config(init_backend["id"], init_backend["config"]) metadata_backend = init_backend["metadata_backend"] - if isinstance(getattr(metadata_backend, "aspect_ratio_bucket_indices", None), dict): - metadata_backend.aspect_ratio_bucket_indices = _coerce_bucket_keys(metadata_backend.aspect_ratio_bucket_indices) if hasattr(metadata_backend, "attach_bucket_report"): metadata_backend.attach_bucket_report(init_backend.get("bucket_report")) if hasattr(metadata_backend, "_mock_children"): diff --git a/simpletuner/helpers/data_backend/webshart.py b/simpletuner/helpers/data_backend/webshart.py index d1d8ab7be..f0685fd36 100644 --- a/simpletuner/helpers/data_backend/webshart.py +++ b/simpletuner/helpers/data_backend/webshart.py @@ -12,9 +12,11 @@ import torch from simpletuner.helpers.data_backend.base import BaseDataBackend +from simpletuner.helpers.data_backend.config.validators import validate_webshart_caption_key from simpletuner.helpers.data_backend.dataset_types import DatasetType, ensure_dataset_type from simpletuner.helpers.data_backend.filters import DatasetFilter from simpletuner.helpers.image_manipulation.load import load_image, load_video +from simpletuner.helpers.prompts import PromptHandler from simpletuner.helpers.training import video_file_extensions from simpletuner.helpers.training.multi_process import should_log @@ -55,9 +57,12 @@ def __init__( compress_cache: bool = False, dataset_type: Union[str, DatasetType] = DatasetType.IMAGE, optimize_captions: bool = False, + caption_key: Optional[Union[str, List[str]]] = None, ): if not source: raise ValueError("source is required for Webshart data backends.") + validate_webshart_caption_key(caption_key) + self.caption_key = caption_key try: import webshart @@ -345,6 +350,14 @@ def get_caption(self, image_path: str) -> Optional[Union[str, List[str], dict]]: sample_ref = self.parse_sample_id(image_path) sample_metadata = self.get_shard_metadata(sample_ref.shard_idx).get(sample_ref.filename, {}) or {} + if self.caption_key is not None: + if sample_metadata.get("json_metadata") is None and sample_metadata.get("json_path"): + reader = self.dataset.open_shard(sample_ref.shard_idx) + payload = reader.read_sample_json(sample_ref.sample_idx) + if payload is not None: + sample_metadata["json_metadata"] = json.loads(payload) + return self._select_caption_keys(sample_metadata) + caption = sample_metadata.get("captions") if caption: if isinstance(caption, dict): @@ -364,6 +377,23 @@ def get_caption(self, image_path: str) -> Optional[Union[str, List[str], dict]]: caption = caption.decode("utf-8") return str(caption).strip() + def _select_caption_keys(self, sample_metadata: dict) -> Optional[Union[str, List[str]]]: + json_metadata = sample_metadata.get("json_metadata") or {} + sources = [json_metadata, sample_metadata] + sources.extend(source.get("captions") for source in list(sources) if isinstance(source, dict)) + keys = [self.caption_key] if isinstance(self.caption_key, str) else self.caption_key + captions = [] + for key in keys: + for source in sources: + if isinstance(source, dict) and key in source: + captions.extend(PromptHandler._normalize_caption_payload(source[key])) + break + if not captions: + return None + if isinstance(self.caption_key, str) and len(captions) == 1: + return captions[0] + return captions + def read(self, identifier: Union[str, Path], as_byteIO: bool = False) -> Any: if self.is_sample_id(identifier): identifier = self.normalize_sample_id(identifier) @@ -501,6 +531,7 @@ def get_instance_representation(self) -> dict: "max_file_size": self.max_file_size, "compress_cache": self.compress_cache, "dataset_type": self.dataset_type.value, + "caption_key": self.caption_key, } @staticmethod @@ -523,6 +554,7 @@ def from_instance_representation(representation: dict) -> "WebshartDataBackend": max_file_size=representation.get("max_file_size", 500 * 1024 * 1024), compress_cache=representation.get("compress_cache", False), dataset_type=representation.get("dataset_type", DatasetType.IMAGE), + caption_key=representation.get("caption_key"), ) def num_shards(self) -> int: diff --git a/simpletuner/helpers/metadata/backends/base.py b/simpletuner/helpers/metadata/backends/base.py index 0cc2dc347..35a8ab953 100644 --- a/simpletuner/helpers/metadata/backends/base.py +++ b/simpletuner/helpers/metadata/backends/base.py @@ -235,7 +235,10 @@ def aspect_ratio_bucket_indices(self): @aspect_ratio_bucket_indices.setter def aspect_ratio_bucket_indices(self, value): - """Set aspect ratio bucket indices with debug tracking.""" + """Keep bucket keys identical during discovery, cache loading, and resume.""" + normalized = {} + for key, samples in value.items(): + normalized.setdefault(str(key), []).extend(samples) if hasattr(self, "_aspect_ratio_bucket_indices"): old_count = sum(len(v) for v in self._aspect_ratio_bucket_indices.values()) new_count = sum(len(v) for v in value.values()) if value else 0 @@ -246,7 +249,7 @@ def aspect_ratio_bucket_indices(self, value): f"Old buckets: {list(self._aspect_ratio_bucket_indices.keys())}, " f"New buckets: {list(value.keys()) if value else []}" ) - self._aspect_ratio_bucket_indices = value + self._aspect_ratio_bucket_indices = normalized def _extract_audio_config(self) -> Dict[str, Any]: if self.dataset_config is None: diff --git a/simpletuner/helpers/metadata/backends/discovery.py b/simpletuner/helpers/metadata/backends/discovery.py index 996ab6d59..4fdac4127 100644 --- a/simpletuner/helpers/metadata/backends/discovery.py +++ b/simpletuner/helpers/metadata/backends/discovery.py @@ -18,19 +18,6 @@ from simpletuner.helpers.training.multi_process import should_log from simpletuner.helpers.training.state_tracker import StateTracker - -def _coerce_bucket_keys_to_float(indices: dict) -> dict: - """Coerce bucket keys from strings to floats (fixes JSON serialization issue).""" - coerced = {} - for key, values in (indices or {}).items(): - try: - coerced_key = float(key) - except (TypeError, ValueError): - coerced_key = key - coerced[coerced_key] = list(values) if not isinstance(values, list) else values - return coerced - - logger = logging.getLogger("DiscoveryMetadataBackend") if should_log(): target_level = os.environ.get("SIMPLETUNER_LOG_LEVEL", "INFO") @@ -340,9 +327,7 @@ def reload_cache(self, set_config: bool = True): except Exception as e: logger.warning(f"Error loading aspect bucket cache, creating new one: {e}") cache_data = {} - # Coerce bucket keys from strings to floats (JSON serialization converts float keys to strings) - loaded_indices = cache_data.get("aspect_ratio_bucket_indices", {}) - self.aspect_ratio_bucket_indices = _coerce_bucket_keys_to_float(loaded_indices) + self.aspect_ratio_bucket_indices = cache_data.get("aspect_ratio_bucket_indices", {}) if set_config: self.config = cache_data.get("config", {}) if self.config != {}: diff --git a/simpletuner/helpers/metadata/backends/huggingface.py b/simpletuner/helpers/metadata/backends/huggingface.py index 302d57f32..2f5fac5d3 100644 --- a/simpletuner/helpers/metadata/backends/huggingface.py +++ b/simpletuner/helpers/metadata/backends/huggingface.py @@ -21,18 +21,6 @@ from simpletuner.helpers.training.state_tracker import StateTracker -def _coerce_bucket_keys_to_float(indices: dict) -> dict: - """Coerce bucket keys from strings to floats (fixes JSON serialization issue).""" - coerced = {} - for key, values in (indices or {}).items(): - try: - coerced_key = float(key) - except (TypeError, ValueError): - coerced_key = key - coerced[coerced_key] = list(values) if not isinstance(values, list) else values - return coerced - - def _dataset_type_value(dataset_type: Any) -> str: return str(getattr(dataset_type, "value", dataset_type)).lower() @@ -385,9 +373,7 @@ def reload_cache(self, set_config: bool = True): except Exception as e: logger.warning(f"Error loading aspect ratio bucket cache, creating new one: {e}") cache_data = {} - # Coerce bucket keys from strings to floats (JSON serialization converts float keys to strings) - loaded_indices = cache_data.get("aspect_ratio_bucket_indices", {}) - self.aspect_ratio_bucket_indices = _coerce_bucket_keys_to_float(loaded_indices) + self.aspect_ratio_bucket_indices = cache_data.get("aspect_ratio_bucket_indices", {}) if set_config: self.config = cache_data.get("config", {}) if self.config != {}: diff --git a/simpletuner/helpers/metadata/backends/parquet.py b/simpletuner/helpers/metadata/backends/parquet.py index 91064d4c1..171762b2f 100644 --- a/simpletuner/helpers/metadata/backends/parquet.py +++ b/simpletuner/helpers/metadata/backends/parquet.py @@ -18,19 +18,6 @@ from simpletuner.helpers.training import audio_file_extensions, image_file_extensions, video_file_extensions from simpletuner.helpers.training.state_tracker import StateTracker - -def _coerce_bucket_keys_to_float(indices: dict) -> dict: - """Coerce bucket keys from strings to floats (fixes JSON serialization issue).""" - coerced = {} - for key, values in (indices or {}).items(): - try: - coerced_key = float(key) - except (TypeError, ValueError): - coerced_key = key - coerced[coerced_key] = list(values) if not isinstance(values, list) else values - return coerced - - logger = logging.getLogger("ParquetMetadataBackend") from simpletuner.helpers.training.multi_process import should_log @@ -278,9 +265,7 @@ def reload_cache(self, set_config: bool = True): except Exception as e: logger.warning(f"Error loading aspect ratio bucket cache, creating new one: {e}") cache_data = {} - # Coerce bucket keys from strings to floats (JSON serialization converts float keys to strings) - loaded_indices = cache_data.get("aspect_ratio_bucket_indices", {}) - self.aspect_ratio_bucket_indices = _coerce_bucket_keys_to_float(loaded_indices) + self.aspect_ratio_bucket_indices = cache_data.get("aspect_ratio_bucket_indices", {}) if set_config: self.config = cache_data.get("config", {}) if self.config != {}: diff --git a/simpletuner/helpers/metadata/backends/webshart.py b/simpletuner/helpers/metadata/backends/webshart.py index 8471a2207..8ed3e6ce1 100644 --- a/simpletuner/helpers/metadata/backends/webshart.py +++ b/simpletuner/helpers/metadata/backends/webshart.py @@ -7,6 +7,7 @@ import time from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext +from hashlib import sha256 from pathlib import Path from typing import Any, Dict, List, Optional, Union @@ -27,17 +28,6 @@ logger.setLevel("ERROR") -def _coerce_bucket_keys_to_float(indices: dict) -> dict: - coerced = {} - for key, values in (indices or {}).items(): - try: - coerced_key = float(key) - except (TypeError, ValueError): - coerced_key = key - coerced[coerced_key] = list(values) if not isinstance(values, list) else values - return coerced - - class WebshartMetadataBackend(MetadataBackend): def __init__( self, @@ -63,6 +53,12 @@ def __init__( repeats: int = 0, max_num_samples: int = None, ): + if not isinstance(data_backend, WebshartDataBackend): + raise ValueError("WebshartMetadataBackend requires WebshartDataBackend") + if data_backend.caption_key is not None: + caption_digest = sha256(json.dumps(data_backend.caption_key).encode("utf-8")).hexdigest()[:16] + cache_file = f"{cache_file}_captions_{caption_digest}" + metadata_file = f"{metadata_file}_captions_{caption_digest}" super().__init__( id=id, instance_data_dir=instance_data_dir, @@ -86,8 +82,6 @@ def __init__( repeats=repeats, max_num_samples=max_num_samples, ) - if not isinstance(data_backend, WebshartDataBackend): - raise ValueError("WebshartMetadataBackend requires WebshartDataBackend") if self.dataset_type not in {DatasetType.IMAGE, DatasetType.VIDEO, DatasetType.CONDITIONING, DatasetType.EVAL}: raise ValueError("WebshartMetadataBackend supports image, video, conditioning, and eval datasets only.") @@ -162,9 +156,7 @@ def reload_cache(self, set_config: bool = True): except Exception as exc: logger.warning("Error loading webshart aspect bucket cache, creating new one: %s", exc) cache_data = {} - self.aspect_ratio_bucket_indices = _coerce_bucket_keys_to_float( - cache_data.get("aspect_ratio_bucket_indices", {}) - ) + self.aspect_ratio_bucket_indices = cache_data.get("aspect_ratio_bucket_indices", {}) self._sync_image_files_with_buckets() if set_config: self.config = cache_data.get("config", {}) @@ -321,6 +313,8 @@ def _metadata_for_entry(self, shard_metadata: dict, filename: str, entry: dict, metadata["original_size"] = (int(width), int(height)) if "captions" in file_metadata: metadata["captions"] = file_metadata["captions"] + if self.data_backend.caption_key is not None: + metadata["captions"] = self.data_backend.get_caption(sample_path) json_metadata = file_metadata.get("json_metadata") or {} if json_metadata: metadata["json_metadata"] = json_metadata @@ -354,7 +348,7 @@ def _prepare_bucket_entry( shard_metadata: dict, entry: dict, sample_path: str, - ) -> tuple[dict, Optional[tuple[float, dict]], Optional[Exception]]: + ) -> tuple[dict, Optional[tuple[str, dict]], Optional[Exception]]: try: filename = str(entry["filename"]) sample_metadata = self._metadata_for_entry(shard_metadata, filename, entry, sample_path) @@ -362,7 +356,7 @@ def _prepare_bucket_entry( except Exception as exc: return {}, None, exc - def _prepare_metadata(self, sample_path: str, sample_metadata: dict) -> Optional[tuple[float, dict]]: + def _prepare_metadata(self, sample_path: str, sample_metadata: dict) -> Optional[tuple[str, dict]]: if not sample_metadata or "original_size" not in sample_metadata: return None if not self.meets_resolution_requirements(image_metadata=sample_metadata): @@ -392,7 +386,7 @@ def _prepare_metadata(self, sample_path: str, sample_metadata: dict) -> Optional ) sample_metadata["bucket_frames"] = rounded_frames else: - bucket_key = round(aspect_ratio, 2) + bucket_key = str(round(aspect_ratio, 2)) return bucket_key, sample_metadata def _entries_for_shard(self, shard_idx: int) -> list[dict]: @@ -526,7 +520,7 @@ def compute_aspect_ratio_bucket_indices(self, ignore_existing_cache: bool = Fals # metadata (e.g. cc12m); get_caption() range-reads those at runtime. # get_shard_metadata returns a flat mapping keyed by member filename. caption_member = Path(str(entry["filename"])).with_suffix(".txt").name - if caption_member not in shard_metadata: + if self.data_backend.caption_key is not None or caption_member not in shard_metadata: statistics["skipped"]["caption_missing"] += 1 continue aspect_ratio_bucket_updates.setdefault(bucket_key, []).append(sample_path) diff --git a/simpletuner/helpers/models/minimaxh3/transformer.py b/simpletuner/helpers/models/minimaxh3/transformer.py index 43f53e617..f3d27a8ea 100644 --- a/simpletuner/helpers/models/minimaxh3/transformer.py +++ b/simpletuner/helpers/models/minimaxh3/transformer.py @@ -69,6 +69,16 @@ _H3_MASKED_CONTEXT_PARALLEL_BACKENDS = frozenset({AttentionBackendName.NATIVE, AttentionBackendName._NATIVE_CUDNN}) +def _linear_compute_dtype(linear: nn.Module) -> torch.dtype: + compute_dtype = getattr(linear, "compute_dtype", None) + if isinstance(compute_dtype, torch.dtype): + return compute_dtype + weight = linear.weight + dequantizer = getattr(weight, "sdnq_dequantizer", None) + result_dtype = getattr(dequantizer, "result_dtype", None) + return result_dtype if isinstance(result_dtype, torch.dtype) else weight.dtype + + class _MiniMaxH3AllGather(torch.autograd.Function): """Gather sequence shards without PyTorch's unsupported NCCL coalesced path.""" @@ -440,7 +450,13 @@ def _infer_minimax_h3_config_from_checkpoint(checkpoint) -> dict[str, Any]: audio_weight = _get_checkpoint_tensor(checkpoint, "audio_proj_in.weight") context_weight = _get_checkpoint_tensor(checkpoint, "context_embedder.weight") q_norm_weight = _get_checkpoint_tensor(checkpoint, "transformer_blocks.0.attn.norm_q.weight") - q_weight = _get_checkpoint_tensor(checkpoint, "transformer_blocks.0.attn.to_q.weight") + if "transformer_blocks.0.attn.to_q.weight" in raw_keys: + q_output_dim = _get_checkpoint_tensor(checkpoint, "transformer_blocks.0.attn.to_q.weight").shape[0] + else: + qkv_weight = _get_checkpoint_tensor(checkpoint, "blocks.0.attn.qkv_proj.weight") + if qkv_weight.shape[0] % 3 != 0: + raise RuntimeError("MiniMax-H3 fused QKV tensor blocks.0.attn.qkv_proj.weight cannot be split into q/k/v") + q_output_dim = qkv_weight.shape[0] // 3 ffn_weight = _get_checkpoint_tensor(checkpoint, "transformer_blocks.0.ff.net.0.proj.weight") has_adaln_curve = "adaln_t_table" in raw_keys adaln_curve_table = _get_checkpoint_tensor(checkpoint, "adaln_t_table") if has_adaln_curve else None @@ -456,7 +472,7 @@ def _infer_minimax_h3_config_from_checkpoint(checkpoint) -> dict[str, Any]: "audio_in_channels": audio_weight.shape[1], "text_dim": context_weight.shape[1], "attention_head_dim": q_norm_weight.shape[0], - "num_attention_heads": q_weight.shape[0] // q_norm_weight.shape[0], + "num_attention_heads": q_output_dim // q_norm_weight.shape[0], "freq_dim": time_in.shape[1] if time_in is not None else 256, "time_embed_hidden_dim": time_in.shape[0] if time_in is not None else 5376, "time_embed_dim": adaln_curve_table.shape[1] if has_adaln_curve else time_out.shape[0], @@ -631,7 +647,7 @@ def forward(self, temb: torch.Tensor) -> tuple[torch.Tensor, ...]: # The activation runs at `temb`'s own precision and only the projection input is aligned to the projection # weight. Every block reads the same `temb`, so early rounding biases every block's modulation coherently. temb = nn.functional.silu(temb) if self.apply_silu else temb - temb = self.linear(temb.to(self.linear.weight.dtype)) + temb = self.linear(temb.to(_linear_compute_dtype(self.linear))) temb = temb.view(-1, 6 * self.hidden_size) return temb.chunk(6, dim=-1) @@ -661,7 +677,7 @@ def forward( ) -> torch.Tensor: # As in `MiniMaxH3AdaLayerNormModulation`: activate at `temb`'s precision, cast to the projection's dtype after. temb = nn.functional.silu(temb) if self.apply_silu else temb - shift, scale = self.linear(temb.to(self.linear.weight.dtype)).chunk(2, dim=-1) + shift, scale = self.linear(temb.to(_linear_compute_dtype(self.linear))).chunk(2, dim=-1) activation_dtype = hidden_states.dtype hidden_states = self.norm(hidden_states) shift = _select_modulation(shift, timestep_indices).to(dtype=activation_dtype) @@ -1522,7 +1538,7 @@ def _time_embedding( temb = blend_flowmap_embeddings(temb, delta_temb, self.flowmap_delta_emb_gate) return temb - dtype = self.time_embedder.linear_1.weight.dtype + dtype = _linear_compute_dtype(self.time_embedder.linear_1) temb = flowmap_timestep_embedding( time_proj=self.time_proj, timestep_embedder=self.time_embedder, @@ -1857,21 +1873,20 @@ def from_single_file( result_dtype=torch_dtype or torch.bfloat16, hadamard_group_size=hadamard_group_size, ) - if len(hadamard_group_sizes) != 1: - raise RuntimeError( - f"MiniMax-H3 ConvRot checkpoint uses multiple Hadamard group sizes: {sorted(hadamard_group_sizes)}" - ) - group_size = hadamard_group_sizes.pop() model.quantization_method = "minimax_h3_comfy_convrot_sdnq" model.quantization_config = { "quant_method": "sdnq_training", "weights_dtype": "int8", "quantized_matmul_dtype": "int8", "use_hadamard": True, - "hadamard_group_size": group_size, "group_size": -1, "source_format": "comfy_minimax_h3_convrot", } + sorted_group_sizes = sorted(hadamard_group_sizes) + if len(sorted_group_sizes) == 1: + model.quantization_config["hadamard_group_size"] = sorted_group_sizes[0] + else: + model.quantization_config["hadamard_group_sizes"] = sorted_group_sizes elif fp8_state_dict: model.quantization_method = "minimax_h3_comfy_fp8" model.quantization_config = { @@ -2022,9 +2037,9 @@ def forward( # mixed-precision (the two patch projections are float32 while `context_embedder` and the block stack are # bfloat16 — see `_keep_in_fp32_modules`), so every input is aligned with its projection's parameter dtype, # mirroring the reference's explicit casts. The text stream sets the dtype of the packed sequence. - video_embeds = self.proj_in(hidden_states.to(self.proj_in.weight.dtype)) - audio_embeds = self.audio_proj_in(audio_hidden_states.to(self.audio_proj_in.weight.dtype)) - text_embeds = self.context_embedder(encoder_hidden_states.to(self.context_embedder.weight.dtype)) + video_embeds = self.proj_in(hidden_states.to(_linear_compute_dtype(self.proj_in))) + audio_embeds = self.audio_proj_in(audio_hidden_states.to(_linear_compute_dtype(self.audio_proj_in))) + text_embeds = self.context_embedder(encoder_hidden_states.to(_linear_compute_dtype(self.context_embedder))) self.token_refiner.gradient_checkpointing = self.gradient_checkpointing text_attention_mask = None if packed_valid_mask is not None: @@ -2375,7 +2390,7 @@ def run_checkpointed_block( # 5. Both heads run over every row, then the rows of each modality are selected. The heads are listed in # `_keep_in_fp32_modules`, so they stay float32 while the block stack runs in the requested `torch_dtype`; # align the activation with their parameter dtype. - hidden_states = self.norm_out(hidden_states, temb, timestep_indices).to(self.proj_out.weight.dtype) + hidden_states = self.norm_out(hidden_states, temb, timestep_indices).to(_linear_compute_dtype(self.proj_out)) video_output = _gather_h3_context_parallel_output(self.proj_out(hidden_states), cp_config, dim=1).index_select( 1, video_indices.to(hidden_states.device) ) diff --git a/simpletuner/helpers/models/z_image/quantized_loading.py b/simpletuner/helpers/models/z_image/quantized_loading.py index c04b53fee..98fee45f5 100644 --- a/simpletuner/helpers/models/z_image/quantized_loading.py +++ b/simpletuner/helpers/models/z_image/quantized_loading.py @@ -200,7 +200,9 @@ def _wrap_convrot_linear( True, -1, ) - _set_module(model, module_name, get_sdnq_wrapper_class(module, forward)) + wrapped_module = get_sdnq_wrapper_class(module, forward) + wrapped_module.compute_dtype = result_dtype + _set_module(model, module_name, wrapped_module) def _validate_quant_metadata(checkpoint, key: str) -> int: diff --git a/simpletuner/helpers/multiaspect/sampler.py b/simpletuner/helpers/multiaspect/sampler.py index 5d1a20475..af4232f2a 100644 --- a/simpletuner/helpers/multiaspect/sampler.py +++ b/simpletuner/helpers/multiaspect/sampler.py @@ -203,14 +203,14 @@ def load_states(self, state_path: str): if isinstance(saved_schedule, dict) and self._saved_schedule_is_restorable(previous_state, state_path): self.metadata_backend.aspect_ratio_bucket_indices = saved_schedule self._val_master_list = sorted(sum(saved_schedule.values(), [])) - self.buckets = previous_state.get("buckets", self.load_buckets()) + self.buckets = [str(bucket) for bucket in previous_state.get("buckets", self.load_buckets())] if "current_bucket" in previous_state: self.current_bucket = previous_state["current_bucket"] self.exhausted_buckets = [] if "exhausted_buckets" in previous_state: self.logger.info(f"Previous checkpoint had {len(previous_state['exhausted_buckets'])} exhausted buckets.") - self.exhausted_buckets = previous_state["exhausted_buckets"] + self.exhausted_buckets = [str(bucket) for bucket in previous_state["exhausted_buckets"]] self.current_epoch = 1 if "current_epoch" in previous_state: self.logger.info(f"Previous checkpoint was on epoch {previous_state['current_epoch']}.") @@ -230,7 +230,7 @@ def load_states(self, state_path: str): self.metadata_backend.seen_images.update(normalized_seen) def load_buckets(self): - return list(self.metadata_backend.aspect_ratio_bucket_indices.keys()) # These keys are a float value, eg. 1.78. + return list(self.metadata_backend.aspect_ratio_bucket_indices.keys()) def retrieve_validation_set(self, batch_size: int): """ @@ -436,36 +436,7 @@ def _reset_buckets(self, raise_exhaustion_signal: bool = True): raise MultiDatasetExhausted() def _get_bucket_images(self, bucket): - """ - Safely retrieve bucket images, trying both original type and type conversion. - - Args: - bucket: The bucket key (could be float or str) - - Returns: - list: List of images in the bucket, or empty list if bucket not found - """ - # Try the original bucket key first - if bucket in self.metadata_backend.aspect_ratio_bucket_indices: - return self.metadata_backend.aspect_ratio_bucket_indices[bucket] - - # Try converting between str and float - try: - if isinstance(bucket, str): - # Try converting str to float - bucket_as_float = float(bucket) - if bucket_as_float in self.metadata_backend.aspect_ratio_bucket_indices: - return self.metadata_backend.aspect_ratio_bucket_indices[bucket_as_float] - elif isinstance(bucket, (float, int)): - # Try converting float/int to str - bucket_as_str = str(bucket) - if bucket_as_str in self.metadata_backend.aspect_ratio_bucket_indices: - return self.metadata_backend.aspect_ratio_bucket_indices[bucket_as_str] - except (ValueError, TypeError): - pass - - # Bucket not found with either type - return [] + return self.metadata_backend.aspect_ratio_bucket_indices.get(str(bucket), []) def _filter_unseen_occurrences(self, images): """Filter consumed positions without collapsing duplicate filepaths.""" @@ -860,7 +831,7 @@ def connect_conditioning_samples(self, samples: tuple): for sample in samples: sample_path: str = sample["image_path"] if self.metadata_backend.instance_data_dir is not None and self.metadata_backend.instance_data_dir != "": - sample_path = sample_path.split(self.metadata_backend.instance_data_dir)[-1] + sample_path = sample_path.split(self.metadata_backend.instance_data_dir)[-1].lstrip("/") # Deterministic selection based on hash of image path and current epoch # This ensures the same image gets the same conditioning dataset within an epoch @@ -916,7 +887,7 @@ def connect_conditioning_samples(self, samples: tuple): for sample in samples: sample_path: str = sample["image_path"] if self.metadata_backend.instance_data_dir is not None and self.metadata_backend.instance_data_dir != "": - sample_path = sample_path.split(self.metadata_backend.instance_data_dir)[-1] + sample_path = sample_path.split(self.metadata_backend.instance_data_dir)[-1].lstrip("/") conditioning_sample = sampler.get_conditioning_sample(sample_path) if conditioning_sample is not None: diff --git a/simpletuner/helpers/publishing/metadata.py b/simpletuner/helpers/publishing/metadata.py index 7f112a548..09c2d161c 100644 --- a/simpletuner/helpers/publishing/metadata.py +++ b/simpletuner/helpers/publishing/metadata.py @@ -736,6 +736,7 @@ def save_model_card( if base_model.count("/") > 1: base_model = f"{model_family}/unknown-model" validation_prompts = _normalize_validation_prompts(validation_prompts) + prompts_by_shortname = dict(zip(validation_shortnames or [], validation_prompts or [])) logger.debug(f"Validating from prompts: {validation_prompts}") assets_folder = os.path.join(repo_folder, "assets") optimizer_config = StateTracker.get_args().optimizer_config @@ -751,7 +752,6 @@ def save_model_card( datasets_str += _dataset_overview_for_model(model, dataset_id, dataset_backend) widget_str = "" idx = 0 - shortname_idx = 0 args = StateTracker.get_args() negative_prompt_text = str(args.validation_negative_prompt) if negative_prompt_text == "": @@ -762,10 +762,23 @@ def save_model_card( audio_sample_rate = model.validation_audio_sample_rate() or 44100 def _add_widget_entries(media, asset_prefix: str): - nonlocal widget_str, idx, shortname_idx, has_video, has_audio - for media_list in media.values() if isinstance(media, dict) else media: + nonlocal widget_str, idx, has_video, has_audio + media_items = media.items() if isinstance(media, dict) else enumerate(media) + for shortname_idx, (shortname, media_list) in enumerate(media_items): if not isinstance(media_list, list): media_list = [media_list] + validation_prompt = "no prompt available" + if validation_prompts is not None: + if isinstance(media, dict) and validation_shortnames is not None: + validation_prompt = prompts_by_shortname.get(shortname, f"prompt not found ({shortname})") + elif shortname_idx < len(validation_prompts): + validation_prompt = validation_prompts[shortname_idx] + else: + validation_prompt = f"prompt not found ({shortname_idx})" + if validation_prompt == "": + validation_prompt = "unconditional (blank prompt)" + else: + validation_prompt = validation_prompt.replace("'", "''") sub_idx = 0 for media_sample in media_list: output_path, media_extension = save_metadata_sample( @@ -779,16 +792,6 @@ def _add_widget_entries(media, asset_prefix: str): has_video = True if media_extension in {"wav", "flac", "mp3", "ogg", "m4a"}: has_audio = True - validation_prompt = "no prompt available" - if validation_prompts is not None: - try: - validation_prompt = validation_prompts[shortname_idx] - except IndexError: - validation_prompt = f"prompt not found ({validation_shortnames[shortname_idx] if validation_shortnames is not None and shortname_idx in validation_shortnames else shortname_idx})" - if validation_prompt == "": - validation_prompt = "unconditional (blank prompt)" - else: - validation_prompt = validation_prompt.replace("'", "''") widget_str += f"\n- text: '{validation_prompt}'" widget_str += "\n parameters:" widget_str += f"\n negative_prompt: '{negative_prompt_text}'" @@ -798,8 +801,6 @@ def _add_widget_entries(media, asset_prefix: str): idx += 1 sub_idx += 1 - shortname_idx += 1 - if has_media: widget_str = "widget:" if images is not None and len(images) > 0: diff --git a/simpletuner/helpers/training/attention_backend.py b/simpletuner/helpers/training/attention_backend.py index 4391238bd..e442706be 100644 --- a/simpletuner/helpers/training/attention_backend.py +++ b/simpletuner/helpers/training/attention_backend.py @@ -5,6 +5,7 @@ import os import subprocess import sys +from contextlib import ExitStack from dataclasses import dataclass from enum import Enum from functools import lru_cache @@ -191,6 +192,13 @@ def _resolve_metal_flash_attention_profile_values( except Exception: continue +_TORCH_SDPA_BACKENDS = { + "_native_math": torch.nn.attention.SDPBackend.MATH, + "_native_flash": torch.nn.attention.SDPBackend.FLASH_ATTENTION, + "_native_efficient": torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION, + "_native_cudnn": torch.nn.attention.SDPBackend.CUDNN_ATTENTION, +} + class AttentionBackendMode(str, Enum): TRAINING = "training" @@ -1460,8 +1468,13 @@ def telemetry_safe_hf_api(*args, **kwargs): cls._disable_diffusers_backend() try: - context = diffusers_attention_backend(backend_enum) - context.__enter__() + with ExitStack() as contexts: + contexts.enter_context(diffusers_attention_backend(backend_enum)) + sdpa_backend = _TORCH_SDPA_BACKENDS.get(_DIFFUSERS_BACKEND_TARGETS.get(backend_key)) + if sdpa_backend is not None: + # Legacy UNet processors call PyTorch SDPA directly, bypassing Diffusers dispatch. + contexts.enter_context(torch.nn.attention.sdpa_kernel(sdpa_backend)) + context = contexts.pop_all() except Exception as exc: message = f"Failed to enable attention backend '{backend_key}': {exc}" logger.error(message) diff --git a/simpletuner/helpers/training/disk_space.py b/simpletuner/helpers/training/disk_space.py index 192dd0f4d..307f082a9 100644 --- a/simpletuner/helpers/training/disk_space.py +++ b/simpletuner/helpers/training/disk_space.py @@ -1,4 +1,4 @@ -"""Disk space monitoring utilities for checkpoint saves.""" +"""Disk space monitoring utilities for checkpoint saves and compilation.""" import logging import re @@ -169,3 +169,16 @@ def check_disk_space( f"{_format_bytes(available)} available, threshold is {threshold_human}." ) logger.info("Disk cleanup script completed. %s now available.", _format_bytes(available)) + + +def check_disk_space_for_config(config: Any, path: str) -> None: + """Apply the configured disk policy to the filesystem containing path.""" + threshold_bytes = parse_size_threshold(getattr(config, "disk_low_threshold", None)) + if threshold_bytes is None: + return + check_disk_space( + output_dir=path, + threshold_bytes=threshold_bytes, + action=DiskLowAction.from_raw(getattr(config, "disk_low_action", None)), + script_path=getattr(config, "disk_low_script", None), + ) diff --git a/simpletuner/helpers/training/dynamo.py b/simpletuner/helpers/training/dynamo.py index 05a6f9147..82a357988 100644 --- a/simpletuner/helpers/training/dynamo.py +++ b/simpletuner/helpers/training/dynamo.py @@ -37,11 +37,30 @@ def dynamo_config_patches(config: Any) -> dict[str, Any]: @contextlib.contextmanager def dynamo_config_context(config: Any) -> Iterator[None]: patches = dynamo_config_patches(config) - if not patches: - yield - return - - with torch._dynamo.config.patch(patches): + with contextlib.ExitStack() as stack: + if patches: + stack.enter_context(torch._dynamo.config.patch(patches)) + if _effective_dynamo_backend(config) == "inductor" and getattr(config, "disk_low_threshold", None) not in ( + None, + "", + "None", + ): + from torch._inductor.runtime.runtime_utils import default_cache_dir + + from simpletuner.helpers.training.disk_space import check_disk_space_for_config + + def check_compile_disk_space(*_args: Any) -> None: + cache = os.environ.get("TORCHINDUCTOR_CACHE_DIR") + if cache is None: + cache = default_cache_dir() + check_disk_space_for_config(config, cache) + triton_cache = os.environ.get("TRITON_CACHE_DIR") + if triton_cache is not None: + check_disk_space_for_config(config, triton_cache) + + check_compile_disk_space() + torch._dynamo.on_compile_start(check_compile_disk_space) + stack.callback(torch._dynamo.callback_handler.remove_start_callback, check_compile_disk_space) yield diff --git a/simpletuner/helpers/training/local_metrics.py b/simpletuner/helpers/training/local_metrics.py index 02b974582..5eb022242 100644 --- a/simpletuner/helpers/training/local_metrics.py +++ b/simpletuner/helpers/training/local_metrics.py @@ -5,6 +5,7 @@ import os from datetime import datetime, timezone from pathlib import Path +from threading import Lock from typing import Any, Optional from accelerate.tracking import GeneralTracker @@ -18,6 +19,7 @@ TIMESTEP_DISTRIBUTION_FILENAME = "timestep_distribution.jsonl" REPORT_FILENAME = "training_report.html" REPORT_REFRESH_INTERVAL = 100 +_REPORT_LOCK = Lock() _CONFIG_KEYS = ( "model_family", @@ -159,28 +161,29 @@ def _json_for_html(payload: Any) -> str: def render_static_report(output_dir: str | os.PathLike[str], max_points: int = 5000) -> Path: - output_path = Path(output_dir).resolve() - package_root = Path(__file__).resolve().parents[2] - template_path = package_root / "templates" / "training_metrics_report.html" - chart_script_path = package_root / "static" / "js" / "training_metrics_chart.js" - chart_style_path = package_root / "static" / "css" / "training_metrics_report.css" - - records = downsample_records(read_metric_records(output_path), max_points=max_points) - payload = { - "run": read_manifest(output_path), - "records": records, - "media": _media_for_report(output_path), - "timesteps": read_timestep_distribution_records(output_path), - } - html = template_path.read_text(encoding="utf-8") - html = html.replace("__SIMPLETUNER_REPORT_CSS__", chart_style_path.read_text(encoding="utf-8")) - html = html.replace("__SIMPLETUNER_REPORT_DATA__", _json_for_html(payload)) - html = html.replace("__SIMPLETUNER_CHART_JS__", chart_script_path.read_text(encoding="utf-8")) - report_path = output_path / REPORT_FILENAME - temporary = report_path.with_suffix(".html.tmp") - temporary.write_text(html, encoding="utf-8") - os.replace(temporary, report_path) - return report_path + with _REPORT_LOCK: + output_path = Path(output_dir).resolve() + package_root = Path(__file__).resolve().parents[2] + template_path = package_root / "templates" / "training_metrics_report.html" + chart_script_path = package_root / "static" / "js" / "training_metrics_chart.js" + chart_style_path = package_root / "static" / "css" / "training_metrics_report.css" + + records = downsample_records(read_metric_records(output_path), max_points=max_points) + payload = { + "run": read_manifest(output_path), + "records": records, + "media": _media_for_report(output_path), + "timesteps": read_timestep_distribution_records(output_path), + } + html = template_path.read_text(encoding="utf-8") + html = html.replace("__SIMPLETUNER_REPORT_CSS__", chart_style_path.read_text(encoding="utf-8")) + html = html.replace("__SIMPLETUNER_REPORT_DATA__", _json_for_html(payload)) + html = html.replace("__SIMPLETUNER_CHART_JS__", chart_script_path.read_text(encoding="utf-8")) + report_path = output_path / REPORT_FILENAME + temporary = report_path.with_suffix(".html.tmp") + temporary.write_text(html, encoding="utf-8") + os.replace(temporary, report_path) + return report_path def record_validation_media( diff --git a/simpletuner/helpers/training/script_runner.py b/simpletuner/helpers/training/script_runner.py index 0a3d7cb59..09c92a9a2 100644 --- a/simpletuner/helpers/training/script_runner.py +++ b/simpletuner/helpers/training/script_runner.py @@ -6,6 +6,7 @@ from concurrent.futures import ThreadPoolExecutor from typing import Callable +from simpletuner.helpers.training.local_metrics import is_local_metrics_enabled, render_static_report from simpletuner.helpers.training.state_tracker import StateTracker from simpletuner.helpers.utils.checkpoint_manager import CheckpointManager @@ -123,4 +124,6 @@ def _resolver(name: str): except ValueError as exc: logger.error("Failed to format external script command: %s", exc) return + if is_local_metrics_enabled(getattr(config, "report_to", None)): + render_static_report(output_dir) submit_script(command) diff --git a/simpletuner/helpers/training/trainer.py b/simpletuner/helpers/training/trainer.py index 12beece93..f191c26da 100644 --- a/simpletuner/helpers/training/trainer.py +++ b/simpletuner/helpers/training/trainer.py @@ -32,11 +32,11 @@ from unittest import mock as unittest_mock import huggingface_hub +import wandb from torch.distributed.fsdp.api import ShardedOptimStateDictConfig, ShardedStateDictConfig from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType from torch.distributed.tensor import DTensor -import wandb from simpletuner.helpers import log_format # noqa from simpletuner.helpers.caching.memory import reclaim_memory from simpletuner.helpers.caching.text_embeds import TextEmbeddingCache @@ -73,6 +73,7 @@ from simpletuner.helpers.training.dynamo import ( apply_checkpointing_cudagraph_compatibility, configure_inductor_wrapper, + dynamo_config_context, install_cudagraph_workarounds, mark_cudagraph_step_begin, ) @@ -1981,58 +1982,59 @@ def _normalize_deepspeed_config(raw_config): def run(self): try: - self._exit_on_signal() - # Initialize essential configurations and schedules - self.init_noise_schedule() - self._exit_on_signal() - self.init_seed() - self._exit_on_signal() - self.init_huggingface_hub() - self._exit_on_signal() - self.init_preprocessing_models() - self._exit_on_signal() - self.init_precision(preprocessing_models_only=True) + with dynamo_config_context(self.config): + self._exit_on_signal() + # Initialize essential configurations and schedules + self.init_noise_schedule() + self._exit_on_signal() + self.init_seed() + self._exit_on_signal() + self.init_huggingface_hub() + self._exit_on_signal() + self.init_preprocessing_models() + self._exit_on_signal() + self.init_precision(preprocessing_models_only=True) - # Core initialization steps with signal checks after each step - self._initialize_components_with_signal_check( - [ - self.init_data_backend, - self.init_validation_prompts, - self.init_unload_text_encoder, - self.init_unload_vae, - self.init_load_base_model, - self.init_delete_model_caches, - self.init_precision, - self.init_controlnet_model, - self.init_tread_model, - self.init_diffusion_blocks_model, - self.init_gligen_layers, - self.init_freeze_models, - self.init_distillation_adapter_modules, - self.init_trainable_peft_adapter, - self.init_diffusion_blocks_trainable_filter, - self.init_lyrics_embedder_training, - ] - ) + # Core initialization steps with signal checks after each step + self._initialize_components_with_signal_check( + [ + self.init_data_backend, + self.init_validation_prompts, + self.init_unload_text_encoder, + self.init_unload_vae, + self.init_load_base_model, + self.init_delete_model_caches, + self.init_precision, + self.init_controlnet_model, + self.init_tread_model, + self.init_diffusion_blocks_model, + self.init_gligen_layers, + self.init_freeze_models, + self.init_distillation_adapter_modules, + self.init_trainable_peft_adapter, + self.init_diffusion_blocks_trainable_filter, + self.init_lyrics_embedder_training, + ] + ) - # Model movement and validation setup - self.move_models(destination="accelerator") - self._exit_on_signal() - self.init_distillation() - self._exit_on_signal() - self.init_validations() - self._exit_on_signal() - AttentionBackendController.apply(self.config, AttentionPhase.EVAL) - self.init_benchmark_base_model() - AttentionBackendController.apply(self.config, AttentionPhase.TRAIN) - self._exit_on_signal() - self.resume_and_prepare() - self._exit_on_signal() - self.init_trackers() - self.run_startup_validation() + # Model movement and validation setup + self.move_models(destination="accelerator") + self._exit_on_signal() + self.init_distillation() + self._exit_on_signal() + self.init_validations() + self._exit_on_signal() + AttentionBackendController.apply(self.config, AttentionPhase.EVAL) + self.init_benchmark_base_model() + AttentionBackendController.apply(self.config, AttentionPhase.TRAIN) + self._exit_on_signal() + self.resume_and_prepare() + self._exit_on_signal() + self.init_trackers() + self.run_startup_validation() - # Start the training process - self.train() + # Start the training process + self.train() except Exception as e: import traceback @@ -6630,32 +6632,12 @@ def checkpoint_state_cleanup(self, output_dir, limit, suffix=None, protected_che def _check_disk_space_before_checkpoint(self) -> None: """Check disk space before saving a checkpoint.""" - threshold_str = getattr(self.config, "disk_low_threshold", None) - if threshold_str in (None, "", "None"): - return - if not self.accelerator.is_main_process: return - from simpletuner.helpers.training.disk_space import DiskLowAction, check_disk_space, parse_size_threshold + from simpletuner.helpers.training.disk_space import check_disk_space_for_config - threshold_bytes = parse_size_threshold(threshold_str) - if threshold_bytes is None: - return - - action = getattr(self.config, "disk_low_action", DiskLowAction.STOP) - if isinstance(action, str): - action = DiskLowAction.from_raw(action) - - script_path = getattr(self.config, "disk_low_script", None) - output_dir = self.config.output_dir - - check_disk_space( - output_dir=output_dir, - threshold_bytes=threshold_bytes, - action=action, - script_path=script_path, - ) + check_disk_space_for_config(self.config, self.config.output_dir) def checkpoint_state_save(self, output_dir, suffix=None): self._check_disk_space_before_checkpoint() diff --git a/simpletuner/helpers/training/validation.py b/simpletuner/helpers/training/validation.py index daa759918..8fe712ad6 100644 --- a/simpletuner/helpers/training/validation.py +++ b/simpletuner/helpers/training/validation.py @@ -5260,10 +5260,17 @@ def disable_ema_for_inference(self): logger.debug("Skipping EMA model restoration for validation, as we are not using EMA.") def _publish_validation_artifacts(self, validation_type: str | None): - if self.publishing_manager is None or not getattr(self.publishing_manager, "configured", False): - return if hasattr(self.accelerator, "is_main_process") and not self.accelerator.is_main_process: return + if self.publishing_manager is None or not getattr(self.publishing_manager, "configured", False): + run_hook_script( + getattr(self.config, "post_upload_script", None), + config=self.config, + local_path=getattr(self.config, "output_dir", None), + remote_path=None, + global_step=self.global_step, + ) + return artifact_root = getattr(self.config, "output_dir", None) if not artifact_root: diff --git a/simpletuner/simpletuner_sdk/server/services/cloud/job_submission.py b/simpletuner/simpletuner_sdk/server/services/cloud/job_submission.py index 421770f54..b6e776444 100644 --- a/simpletuner/simpletuner_sdk/server/services/cloud/job_submission.py +++ b/simpletuner/simpletuner_sdk/server/services/cloud/job_submission.py @@ -49,6 +49,16 @@ class SubmissionContext: reservation_id: Optional[str] = None quota_warnings: List[str] = field(default_factory=list) + @property + def tracker_metadata(self) -> Dict[str, str]: + metadata = {} + project = self.config.get("--tracker_project_name") or self.config.get("tracker_project_name") + if project: + metadata["tracker_project_name"] = project + if self.tracker_run_name: + metadata["tracker_run_name"] = self.tracker_run_name + return metadata + @dataclass class SubmissionResult: @@ -206,8 +216,7 @@ async def submit( unified_job.output_url = f"https://huggingface.co/{hub_model_id}" if snapshot_metadata: unified_job.metadata["snapshot"] = snapshot_metadata - if ctx.tracker_run_name: - unified_job.metadata["tracker_run_name"] = ctx.tracker_run_name + unified_job.metadata.update(ctx.tracker_metadata) if ctx.hardware_profile: unified_job.metadata["hardware_profile"] = ctx.hardware_profile @@ -296,8 +305,7 @@ async def _create_upload_job( metadata: Dict[str, Any] = {"upload_id": job_id} if snapshot_metadata: metadata["snapshot"] = snapshot_metadata - if ctx.tracker_run_name: - metadata["tracker_run_name"] = ctx.tracker_run_name + metadata.update(ctx.tracker_metadata) if ctx.hardware_profile: metadata["hardware_profile"] = ctx.hardware_profile @@ -327,8 +335,7 @@ async def _update_upload_job_from_provider( metadata.setdefault("prediction_id", cloud_job.job_id) if snapshot_metadata: metadata["snapshot"] = snapshot_metadata - if ctx.tracker_run_name: - metadata["tracker_run_name"] = ctx.tracker_run_name + metadata.update(ctx.tracker_metadata) if ctx.upload_id: metadata["upload_id"] = ctx.upload_id if ctx.hardware_profile: diff --git a/simpletuner/simpletuner_sdk/server/services/configs_service.py b/simpletuner/simpletuner_sdk/server/services/configs_service.py index 81fe5c05c..78adc8acf 100644 --- a/simpletuner/simpletuner_sdk/server/services/configs_service.py +++ b/simpletuner/simpletuner_sdk/server/services/configs_service.py @@ -1461,6 +1461,7 @@ def normalize_form_to_config( "--publishing_config", "--lycoris_config", "--user_prompt_library", + "--validation_adapter_config", "--peft_lora_target_modules", } diff --git a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/advanced.py b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/advanced.py index 1ad3479c7..1ba3fd025 100644 --- a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/advanced.py +++ b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/advanced.py @@ -2724,7 +2724,7 @@ def register_advanced_fields(registry: "FieldRegistry") -> None: section="checkpointing", subsection="advanced", default_value=None, - help_text="Minimum free disk space required before checkpoint saves (e.g., '100G', '50M'). Leave empty to disable.", + help_text="Minimum free disk space required before checkpoint saves and TorchInductor compilation (e.g., '100G', '50M'). Leave empty to disable.", tooltip="Training will pause or stop if available space drops below this threshold.", importance=ImportanceLevel.ADVANCED, group="disk_low_space", diff --git a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/publishing.py b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/publishing.py index 4d4fd7c1a..94daf8412 100644 --- a/simpletuner/simpletuner_sdk/server/services/field_registry/sections/publishing.py +++ b/simpletuner/simpletuner_sdk/server/services/field_registry/sections/publishing.py @@ -79,7 +79,7 @@ def register_publishing_fields(registry: "FieldRegistry") -> None: section="publishing_controls", default_value=None, placeholder="/path/to/hook.sh {remote_checkpoint_path}", - help_text="Optional executable to run after each publishing provider and Hugging Face Hub upload finishes.", + help_text="Optional executable run after each publishing provider or Hugging Face Hub upload, and after completed built-in validation when no publishing provider is configured.", tooltip="Supports placeholders like {remote_checkpoint_path}, {local_checkpoint_path}, {global_step}, {tracker_run_name}, {tracker_project_name}, {model_family}, {huggingface_path}. Runs asynchronously.", importance=ImportanceLevel.ADVANCED, order=6, diff --git a/simpletuner/static/css/training_metrics_report.css b/simpletuner/static/css/training_metrics_report.css index 1bda806ff..40e66aa80 100644 --- a/simpletuner/static/css/training_metrics_report.css +++ b/simpletuner/static/css/training_metrics_report.css @@ -201,6 +201,7 @@ input { } .media-timeline { + min-width: 0; display: flex; gap: 6px; margin-bottom: 12px; @@ -208,7 +209,16 @@ input { padding-bottom: 4px; } +#media-timeline { + position: sticky; + top: 0; + z-index: 3; + padding: 8px 0; + background: var(--report-bg); +} + .media-step { + min-height: 44px; flex: 0 0 auto; padding: 6px 9px; border: 1px solid var(--report-border); @@ -288,7 +298,7 @@ input { inset: 0; z-index: 10; display: grid; - grid-template-rows: minmax(0, 1fr) auto auto; + grid-template-rows: minmax(0, 1fr) auto auto auto; gap: 12px; padding: 16px; background: rgba(5, 7, 11, 0.95); @@ -319,6 +329,13 @@ input { max-height: none; } +#media-lightbox-steps { + width: 100%; + max-width: 600px; + justify-self: center; + margin-bottom: 0; +} + .media-lightbox-toolbar { display: flex; justify-content: center; @@ -326,8 +343,8 @@ input { } .media-lightbox-toolbar button { - min-width: 36px; - min-height: 36px; + min-width: 44px; + min-height: 44px; border: 1px solid var(--report-border); background: var(--report-panel); color: var(--report-text); diff --git a/simpletuner/static/js/dataloader-section-component.js b/simpletuner/static/js/dataloader-section-component.js index 094c5993d..eff3b09ca 100644 --- a/simpletuner/static/js/dataloader-section-component.js +++ b/simpletuner/static/js/dataloader-section-component.js @@ -1348,6 +1348,16 @@ function dataloaderSectionComponent() { trainer.markDatasetsDirty(); } }, + setWebshartCaptionKeys(dataset, value) { + const keys = value.split('\n').map(key => key.trim()).filter(Boolean); + if (keys.length) { + dataset.webshart = dataset.webshart || {}; + dataset.webshart.caption_key = keys.length === 1 ? keys[0] : keys; + } else if (dataset.webshart) { + delete dataset.webshart.caption_key; + } + this.markAsUnsaved(); + }, onStorageBackendChange(dataset) { if (!dataset) { return; diff --git a/simpletuner/static/js/modules/cloud/index.js b/simpletuner/static/js/modules/cloud/index.js index 772c1e683..0781cc259 100644 --- a/simpletuner/static/js/modules/cloud/index.js +++ b/simpletuner/static/js/modules/cloud/index.js @@ -612,7 +612,7 @@ if (!window.cloudDashboardComponent) { filtered = filtered.filter(j => { const jobId = (j.job_id || '').toLowerCase(); const configName = (j.config_name || '').toLowerCase(); - const trackerName = (j.metadata?.tracker_run_name || '').toLowerCase(); + const trackerName = this.jobDisplayName(j).toLowerCase(); const status = (j.status || '').toLowerCase(); const provider = (j.provider || '').toLowerCase(); return jobId.includes(query) || configName.includes(query) || diff --git a/simpletuner/static/js/modules/cloud/jobs.js b/simpletuner/static/js/modules/cloud/jobs.js index 9a1e09206..3ce34fe34 100644 --- a/simpletuner/static/js/modules/cloud/jobs.js +++ b/simpletuner/static/js/modules/cloud/jobs.js @@ -5,6 +5,25 @@ */ window.cloudJobMethods = { + async continueJob(job) { + const trainer = Alpine.store('trainer'); + if (trainer.continuingJob || job.job_type !== 'local' || job.status !== 'failed') return; + trainer.continuingJob = job.job_id; + try { + const switched = await trainer.switchEnvironment(job.metadata?.env_name || job.config_name); + if (!switched) return; + trainer.switchTab('basic'); + await htmx.ajax('GET', '/web/trainer/tabs/basic', { target: '#tab-content', swap: 'innerHTML' }); + await Alpine.nextTick(); + document.getElementById('runBtn').click(); + } catch (error) { + console.error('Failed to continue job:', error); + window.showToast('Failed to continue job', 'error'); + } finally { + trainer.continuingJob = null; + } + }, + async loadJobs(syncActive = false) { if (!syncActive) { this.jobsLoading = true; diff --git a/simpletuner/static/js/modules/cloud/utilities.js b/simpletuner/static/js/modules/cloud/utilities.js index 9f189e42c..75e710d40 100644 --- a/simpletuner/static/js/modules/cloud/utilities.js +++ b/simpletuner/static/js/modules/cloud/utilities.js @@ -46,7 +46,12 @@ window.cloudUtilityMethods = { }, jobDisplayName(job) { - return job.metadata?.tracker_run_name || job.config_name || job.job_id.substring(0, 8); + const metadata = job.metadata || {}; + const config = metadata.runtime_config || {}; + const project = metadata.tracker_project_name || config['--tracker_project_name'] || config.tracker_project_name; + const run = metadata.tracker_run_name || config['--tracker_run_name'] || config.tracker_run_name || + metadata.run_name || job.config_name || job.job_id.substring(0, 8); + return project ? `${project} / ${run}` : run; }, formatDuration(seconds) { @@ -204,7 +209,7 @@ window.cloudComputedProperties = { filtered = filtered.filter(j => { const jobId = (j.job_id || '').toLowerCase(); const configName = (j.config_name || '').toLowerCase(); - const trackerName = (j.metadata?.tracker_run_name || '').toLowerCase(); + const trackerName = window.cloudUtilityMethods.jobDisplayName(j).toLowerCase(); const status = (j.status || '').toLowerCase(); const provider = (j.provider || '').toLowerCase(); diff --git a/simpletuner/static/js/training_metrics_chart.js b/simpletuner/static/js/training_metrics_chart.js index 853e884f9..d2d85b93c 100644 --- a/simpletuner/static/js/training_metrics_chart.js +++ b/simpletuner/static/js/training_metrics_chart.js @@ -65,8 +65,8 @@ function formatMetricValue(value) { if (typeof value !== 'number' || !Number.isFinite(value)) return '—'; const absolute = Math.abs(value); - if (absolute !== 0 && (absolute >= 10000 || absolute < 0.001)) return value.toExponential(3); - return value.toLocaleString(undefined, { maximumFractionDigits: 6 }); + if (absolute !== 0 && (absolute >= 10000 || absolute < 0.001)) return value.toExponential(2); + return value.toLocaleString(undefined, { maximumFractionDigits: 2 }); } function timestampMs(record) { diff --git a/simpletuner/templates/components/dataloader/sections/storage.html b/simpletuner/templates/components/dataloader/sections/storage.html index 8b7d7457a..021c3653e 100644 --- a/simpletuner/templates/components/dataloader/sections/storage.html +++ b/simpletuner/templates/components/dataloader/sections/storage.html @@ -259,6 +259,9 @@