From bbb20177cb57ae5e913b63134c63cedffebf91b1 Mon Sep 17 00:00:00 2001
From: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>
Date: Wed, 20 May 2026 12:32:53 +0100
Subject: [PATCH 1/3] feat(embed): add optimizer backend selection

Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>
---
 src/nemotron/recipes/embed/README.md          |   6 +-
 .../embed/stage2_finetune/biencoder_base.yaml |  23 +-
 .../embed/stage2_finetune/config/default.yaml |   8 +-
 .../embed/stage2_finetune/pyproject.toml      |   9 +-
 .../recipes/embed/stage2_finetune/train.py    | 109 +++++-
 .../recipes/embed/stage2_finetune/uv.lock     | 323 +++++++++++++-----
 tests/recipes/embed/integration/conftest.py   |   2 +-
 .../embed/test_finetune_optimizer_backends.py |  60 ++++
 8 files changed, 424 insertions(+), 116 deletions(-)
 create mode 100644 tests/recipes/embed/test_finetune_optimizer_backends.py

diff --git a/src/nemotron/recipes/embed/README.md b/src/nemotron/recipes/embed/README.md
index 8cfe39d09..b8f7357d5 100644
--- a/src/nemotron/recipes/embed/README.md
+++ b/src/nemotron/recipes/embed/README.md
@@ -356,7 +356,7 @@ entity = "my-team"
 # Local Docker execution profile
 [local-docker]
 executor = "docker"
-container_image = "nvcr.io/nvidia/pytorch:25.01-py3"
+container_image = "nvcr.io/nvidia/nemo-automodel:26.04"
 runtime = "nvidia"  # Enable GPU passthrough
 ipc_mode = "host"
 shm_size = "16g"
@@ -371,7 +371,7 @@ executor = "slurm"
 account = "my-account"
 partition = "interactive"
 batch_partition = "batch"
-container_image = "nvcr.io/nvidia/pytorch:25.01-py3"
+container_image = "nvcr.io/nvidia/nemo-automodel:26.04"
 tunnel = "ssh"
 host = "cluster.example.com"
 user = "username"
@@ -451,6 +451,8 @@ base_model: nvidia/llama-nemotron-embed-1b-v2
 num_epochs: 3
 global_batch_size: 128
 learning_rate: 1.0e-5
+optimizer_backend: auto        # FusedAdam in Automodel container, FlashAdamW fallback
+flash_adamw_master_weight_bits: 32
 query_max_length: 512          # Max query tokens (check your base model's max sequence length)
 passage_max_length: 512        # Max passage tokens (check your base model's max sequence length)
 # attn_implementation: null    # Auto-detects flash_attention_2 if available, else sdpa
diff --git a/src/nemotron/recipes/embed/stage2_finetune/biencoder_base.yaml b/src/nemotron/recipes/embed/stage2_finetune/biencoder_base.yaml
index be5dcaa50..cd9fe67f5 100644
--- a/src/nemotron/recipes/embed/stage2_finetune/biencoder_base.yaml
+++ b/src/nemotron/recipes/embed/stage2_finetune/biencoder_base.yaml
@@ -30,22 +30,17 @@ dist_env:
   timeout_minutes: 1
 
 model:
-  _target_: nemo_automodel.components.models.biencoder.NeMoAutoModelBiencoder.from_pretrained
+  _target_: nemo_automodel._transformers.auto_model.NeMoAutoModelBiEncoder.from_pretrained
   pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
-  share_encoder: true
-  add_linear_pooler: false
-  out_dimension: 768
-  do_gradient_checkpointing: false
-  train_n_passages: 5
-  eval_negative_size: 4
   attn_implementation: sdpa  # Use PyTorch SDPA as fallback (works with/without flash-attn)
   pooling: avg
   l2_normalize: true
-  t: 0.02
   use_liger_kernel: true
   use_sdpa_patching: true
   torch_dtype: bfloat16
 
+temperature: 0.02
+
 tokenizer:
   _target_: nemo_automodel._transformers.auto_tokenizer.NeMoAutoTokenizer.from_pretrained
   pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
@@ -56,13 +51,14 @@ dataloader:
     _target_: nemo_automodel.components.datasets.llm.make_retrieval_dataset
     data_dir_list:
       - training_datasets/nqsh_shuffled_50k.json
+    model_type: bi_encoder
     data_type: train
-    train_n_passages: 5
+    n_passages: 5
     eval_negative_size: 4
     seed: 42
     do_shuffle: true
   collate_fn:
-    _target_: nemo_automodel.components.datasets.llm.RetrievalBiencoderCollator
+    _target_: nemo_automodel.components.datasets.llm.BiEncoderCollator
     q_max_len: 512
     p_max_len: 512
     query_prefix: "query:"
@@ -72,11 +68,12 @@ dataloader:
   num_workers: 0
 
 optimizer:
-  _target_: torch.optim.AdamW
+  _target_: transformer_engine.pytorch.optimizers.fused_adam.FusedAdam
   lr: 5.0e-6
   weight_decay: 0.01
-  betas: [0.9, 0.999]
-  eps: 1.0e-8
+  adam_w_mode: true
+  bias_correction: true
+  master_weights: true
 
 lr_scheduler:
   lr_warmup_steps: 1
diff --git a/src/nemotron/recipes/embed/stage2_finetune/config/default.yaml b/src/nemotron/recipes/embed/stage2_finetune/config/default.yaml
index 62be5e246..32ff77660 100644
--- a/src/nemotron/recipes/embed/stage2_finetune/config/default.yaml
+++ b/src/nemotron/recipes/embed/stage2_finetune/config/default.yaml
@@ -14,7 +14,7 @@
 run:
   env:
     # Container image for remote execution (Docker/Slurm)
-    container: nvcr.io/nvidia/pytorch:25.12-py3
+    container: nvcr.io/nvidia/nemo-automodel:26.04
 
 # Base embedding model to fine-tune
 base_model: nvidia/llama-nemotron-embed-1b-v2
@@ -34,6 +34,12 @@ lr_warmup_steps: 5
 lr_decay_style: cosine  # cosine maintains higher LR longer than linear
 weight_decay: 0.01
 
+# Optimizer
+# auto uses Transformer Engine FusedAdam when available in the Automodel
+# container, otherwise FlashAdamW with bf16 model parameters.
+optimizer_backend: auto
+flash_adamw_master_weight_bits: 32
+
 # Model architecture
 # attn_implementation: null  # Auto-detects: flash_attention_2 if available, else sdpa
 train_n_passages: 5  # 1 positive + 4 negatives
diff --git a/src/nemotron/recipes/embed/stage2_finetune/pyproject.toml b/src/nemotron/recipes/embed/stage2_finetune/pyproject.toml
index 7458a1dcb..731324e4b 100644
--- a/src/nemotron/recipes/embed/stage2_finetune/pyproject.toml
+++ b/src/nemotron/recipes/embed/stage2_finetune/pyproject.toml
@@ -3,20 +3,21 @@ name = "recipe-runner"
 version = "0.0.0"
 requires-python = ">=3.12,<3.13"
 dependencies = [
-  "nemo-automodel @ git+https://github.com/NVIDIA-NeMo/Automodel.git@ecd7cb4297ae06990c94b7581f34c03ce8ed5488",
+  "nemo-automodel==0.4.0",
   # Needed by Nemotron library (added via --with flag)
   "omegaconf>=2.3.0",
   "pydantic-settings>=2.0.0",
   # Local GPU torch wheel, sourced from PyTorch cu129 index on Linux below
   "torch>=2.10.0",
-  # Pin transformers <5.2 — nemo-automodel's check_model_inputs decorator
-  # is a ContextDecorator in 5.2.0 which breaks biencoder forward()
-  "transformers>=5.0,<5.2",
 ]
 
 # Note: exclude-dependencies are injected dynamically by run_uv.py for Docker/Slurm
 # This base config works for local mode where all packages need to be installed
 
+[tool.uv]
+# Override automodel's transformers pin to use 5.5.x
+override-dependencies = ["transformers>=5.5.0,<5.6.0"]
+
 # Pin Linux torch installs to cu129. CUDA 13 drivers can run cu12.x wheels via
 # NVIDIA driver backward compatibility, while CUDA 12.9 hosts cannot run cu130.
 [tool.uv.sources]
diff --git a/src/nemotron/recipes/embed/stage2_finetune/train.py b/src/nemotron/recipes/embed/stage2_finetune/train.py
index c54bafb2b..edc03a7f0 100644
--- a/src/nemotron/recipes/embed/stage2_finetune/train.py
+++ b/src/nemotron/recipes/embed/stage2_finetune/train.py
@@ -4,8 +4,8 @@
 # schema = "1"
 # docs = "https://raw.githubusercontent.com/NVIDIA-NeMo/Nemotron/main/docs/runspec/v1/spec.md"
 # name = "embed/finetune"
-# image = "nvcr.io/nvidia/pytorch:25.12-py3"
-# setup = "PyTorch pre-installed. Stage dependencies resolved via UV at runtime."
+# image = "nvcr.io/nvidia/nemo-automodel:26.04"
+# setup = "NeMo Automodel pre-installed. Stage dependencies resolved via UV at runtime."
 #
 # [tool.runspec.run]
 # launch = "torchrun"
@@ -50,11 +50,12 @@
 
 from __future__ import annotations
 
+import importlib
 import json
 import os
 import sys
 from pathlib import Path
-from typing import Literal
+from typing import Any, Literal
 
 from pydantic import ConfigDict, Field
 
@@ -107,6 +108,14 @@ class FinetuneConfig(RecipeSettings):
         description="LR decay schedule (cosine, linear).",
     )
     weight_decay: float = Field(default=0.01, ge=0, description="Weight decay for optimizer.")
+    optimizer_backend: Literal["auto", "fused_adam", "flash_adamw"] = Field(
+        default="auto",
+        description="Optimizer backend. 'auto' uses FusedAdam when available, otherwise FlashAdamW.",
+    )
+    flash_adamw_master_weight_bits: Literal[24, 32] = Field(
+        default=32,
+        description="Effective master-weight precision for FlashAdamW when Transformer Engine is unavailable.",
+    )
 
     # Model architecture
     attn_implementation: Literal["sdpa", "flash_attention_2", "eager"] | None = Field(
@@ -213,6 +222,74 @@ def _auto_scale_hyperparams(
     return global_batch_size, num_epochs, checkpoint_every_steps, val_every_steps
 
 
+def _can_import_fused_adam() -> tuple[bool, str | None]:
+    """Return whether Transformer Engine FusedAdam is importable."""
+    try:
+        importlib.import_module("transformer_engine.pytorch.optimizers.fused_adam")
+    except Exception as e:
+        return False, str(e)
+    return True, None
+
+
+def _can_import_flash_adamw() -> tuple[bool, str | None]:
+    """Return whether FlashAdamW is importable."""
+    try:
+        importlib.import_module("flashoptim")
+    except Exception as e:
+        return False, str(e)
+    return True, None
+
+
+def _load_automodel_config(cfg: FinetuneConfig, config_node_cls: type) -> tuple[Any, str]:
+    """Load Automodel YAML after choosing an optimizer that is importable here."""
+    import yaml
+
+    base_config_path = STAGE_PATH / "biencoder_base.yaml"
+    with open(base_config_path) as f:
+        raw_config = yaml.safe_load(f)
+
+    te_available, te_error = _can_import_fused_adam()
+    flash_available, flash_error = _can_import_flash_adamw()
+    optimizer_backend = cfg.optimizer_backend
+    if optimizer_backend == "auto":
+        optimizer_backend = "fused_adam" if te_available else "flash_adamw"
+
+    if optimizer_backend == "fused_adam":
+        if not te_available:
+            print("Error: optimizer_backend=fused_adam requires Transformer Engine.", file=sys.stderr)
+            if te_error:
+                print(f"  Import error: {te_error}", file=sys.stderr)
+            print(
+                "  Use optimizer_backend=flash_adamw for local runs without Transformer Engine.",
+                file=sys.stderr,
+            )
+            sys.exit(1)
+    elif optimizer_backend == "flash_adamw":
+        if not flash_available:
+            print("Error: optimizer_backend=flash_adamw requires flashoptim.", file=sys.stderr)
+            if flash_error:
+                print(f"  Import error: {flash_error}", file=sys.stderr)
+            print(
+                "  Install flashoptim, or run in an environment with Transformer Engine FusedAdam.",
+                file=sys.stderr,
+            )
+            sys.exit(1)
+        raw_config["optimizer"] = {
+            "_target_": "flashoptim.FlashAdamW",
+            "lr": raw_config.get("optimizer", {}).get("lr", cfg.learning_rate),
+            "weight_decay": raw_config.get("optimizer", {}).get("weight_decay", cfg.weight_decay),
+            "betas": [0.9, 0.999],
+            "eps": 1.0e-8,
+            "quantize": False,
+            "compress_state_dict": False,
+            "master_weight_bits": cfg.flash_adamw_master_weight_bits,
+            "fused": True,
+        }
+        raw_config.setdefault("model", {})["torch_dtype"] = "bfloat16"
+
+    return config_node_cls(raw_config), optimizer_backend
+
+
 def run_finetune(cfg: FinetuneConfig) -> Path:
     """Run embedding model fine-tuning using nemo-automodel.
 
@@ -278,17 +355,25 @@ def run_finetune(cfg: FinetuneConfig) -> Path:
 
     # Import nemo-automodel components
     try:
-        from nemo_automodel.components.config.loader import load_yaml_config
-        from nemo_automodel.recipes.biencoder import TrainBiencoderRecipe
+        from nemo_automodel.components.config.loader import ConfigNode
+        from nemo_automodel.recipes.retrieval import TrainBiEncoderRecipe
     except ImportError as e:
         print("Error: Failed to import nemo-automodel. Is it installed?", file=sys.stderr)
         print("  Install with: pip install nemo-automodel", file=sys.stderr)
         print(f"  Error: {e}", file=sys.stderr)
         sys.exit(1)
 
-    # Load base config from nemo-automodel defaults
-    base_config_path = STAGE_PATH / "biencoder_base.yaml"
-    automodel_cfg = load_yaml_config(str(base_config_path))
+    # Load base config from nemo-automodel defaults. ConfigNode resolves _target_
+    # imports during construction, so optimizer selection must happen on raw YAML.
+    automodel_cfg, optimizer_backend = _load_automodel_config(cfg, ConfigNode)
+    optimizer_detail = optimizer_backend
+    if optimizer_backend == "flash_adamw":
+        optimizer_detail = (
+            f"{optimizer_backend} "
+            f"(bf16 model, {cfg.flash_adamw_master_weight_bits}-bit master weights)"
+        )
+    print(f"Optimizer:      {optimizer_detail}")
+    print()
 
     # Apply overrides from our config
     # Model settings
@@ -308,7 +393,7 @@ def run_finetune(cfg: FinetuneConfig) -> Path:
 
     # Data settings
     automodel_cfg.dataloader.dataset.data_dir_list = [str(cfg.train_data_path)]
-    automodel_cfg.dataloader.dataset.train_n_passages = cfg.train_n_passages
+    automodel_cfg.dataloader.dataset.n_passages = cfg.train_n_passages
     automodel_cfg.dataloader.collate_fn.q_max_len = cfg.query_max_length
     automodel_cfg.dataloader.collate_fn.p_max_len = cfg.passage_max_length
     automodel_cfg.dataloader.collate_fn.query_prefix = cfg.query_prefix
@@ -330,13 +415,13 @@ def run_finetune(cfg: FinetuneConfig) -> Path:
     # Model architecture
     automodel_cfg.model.pooling = cfg.pooling
     automodel_cfg.model.l2_normalize = cfg.l2_normalize
-    automodel_cfg.model.t = cfg.temperature
+    automodel_cfg.temperature = cfg.temperature
 
     # Checkpoint settings
     automodel_cfg.checkpoint.checkpoint_dir = str(cfg.checkpoint_dir)
 
-    # Create and run the biencoder recipe
-    recipe = TrainBiencoderRecipe(automodel_cfg)
+    # Create and run the bi-encoder recipe
+    recipe = TrainBiEncoderRecipe(automodel_cfg)
     recipe.setup()
     recipe.run_train_validation_loop()
 
diff --git a/src/nemotron/recipes/embed/stage2_finetune/uv.lock b/src/nemotron/recipes/embed/stage2_finetune/uv.lock
index 3635edbff..8f5ee0168 100644
--- a/src/nemotron/recipes/embed/stage2_finetune/uv.lock
+++ b/src/nemotron/recipes/embed/stage2_finetune/uv.lock
@@ -8,6 +8,9 @@ resolution-markers = [
     "sys_platform == 'darwin'",
 ]
 
+[manifest]
+overrides = [{ name = "transformers", specifier = ">=5.5.0,<5.6.0" }]
+
 [[package]]
 name = "aiohappyeyeballs"
 version = "2.6.1"
@@ -358,26 +361,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/b0/e5/247d094108e42ac26363ab8dc57f168840cf7c05774b40ffeb0d78868fcc/datasets-4.8.4-py3-none-any.whl", hash = "sha256:cdc8bee4698e549d78bf1fed6aea2eebc760b22b084f07e6fc020c6577a6ce6d", size = 526991, upload-time = "2026-03-23T14:21:15.89Z" },
 ]
 
-[[package]]
-name = "diffusers"
-version = "0.37.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "filelock" },
-    { name = "httpx" },
-    { name = "huggingface-hub" },
-    { name = "importlib-metadata" },
-    { name = "numpy" },
-    { name = "pillow" },
-    { name = "regex" },
-    { name = "requests" },
-    { name = "safetensors" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/46/5c/f4c2eb8d481fe8784a7e2331fbaab820079c06676185fa6d2177b386d590/diffusers-0.37.1.tar.gz", hash = "sha256:2346c21f77f835f273b7aacbaada1c34a596a3a2cc6ddc99d149efcd0ec298fa", size = 4135139, upload-time = "2026-03-25T08:04:04.515Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/9c/dd/51c38785ce5e1c287b5ad17ba550edaaaffce0deb0da4857019c6700fbaf/diffusers-0.37.1-py3-none-any.whl", hash = "sha256:0537c0b28cb53cf39d6195489bcf8f833986df556c10f5e28ab7427b86fc8b90", size = 5001536, upload-time = "2026-03-25T08:04:02.385Z" },
-]
-
 [[package]]
 name = "dill"
 version = "0.4.1"
@@ -435,6 +418,21 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/a4/a5/842ae8f0c08b61d6484b52f99a03510a3a72d23141942d216ebe81fefbce/filelock-3.25.2-py3-none-any.whl", hash = "sha256:ca8afb0da15f229774c9ad1b455ed96e85a81373065fb10446672f64444ddf70", size = 26759, upload-time = "2026-03-11T20:45:37.437Z" },
 ]
 
+[[package]]
+name = "flashoptim"
+version = "0.1.4"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "torch", version = "2.10.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin'" },
+    { name = "torch", version = "2.10.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform != 'darwin' and sys_platform != 'linux'" },
+    { name = "torch", version = "2.10.0+cu129", source = { registry = "https://download.pytorch.org/whl/cu129" }, marker = "sys_platform == 'linux'" },
+    { name = "triton", marker = "sys_platform == 'linux'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/a3/ad/9bf3ec596650e23371a1515014de3d250046e18de8226310ebd6256677a9/flashoptim-0.1.4.tar.gz", hash = "sha256:8d2d9f2608bef51005ba827b8b63af17189979cc0301c542aad2610ab1780718", size = 99938, upload-time = "2026-04-17T18:17:49.268Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/08/2e/8465a97b6d574f50ecb98b76b2a57b8306b1834478f2097abf57a77d2bcf/flashoptim-0.1.4-py3-none-any.whl", hash = "sha256:8a4a3f2528fbda419d4f4dd0c9debb3de22bd0a45969bee2eb5a58185d3bd451", size = 46638, upload-time = "2026-04-17T18:17:51.047Z" },
+]
+
 [[package]]
 name = "flask"
 version = "3.1.3"
@@ -521,18 +519,6 @@ http = [
     { name = "aiohttp" },
 ]
 
-[[package]]
-name = "ftfy"
-version = "6.3.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "wcwidth" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/a5/d3/8650919bc3c7c6e90ee3fa7fd618bf373cbbe55dff043bd67353dbb20cd8/ftfy-6.3.1.tar.gz", hash = "sha256:9b3c3d90f84fb267fe64d375a07b7f8912d817cf86009ae134aa03e1819506ec", size = 308927, upload-time = "2024-10-26T00:50:35.149Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/ab/6e/81d47999aebc1b155f81eca4477a616a70f238a2549848c38983f3c22a82/ftfy-6.3.1-py3-none-any.whl", hash = "sha256:7c70eb532015cd2f9adb53f101fb6c7945988d023a085d127d1573dc49dd0083", size = 44821, upload-time = "2024-10-26T00:50:33.425Z" },
-]
-
 [[package]]
 name = "gitdb"
 version = "4.0.12"
@@ -725,20 +711,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/0e/61/66938bbb5fc52dbdf84594873d5b51fb1f7c7794e9c0f5bd885f30bc507b/idna-3.11-py3-none-any.whl", hash = "sha256:771a87f49d9defaf64091e6e6fe9c18d4833f140bd19464795bc32d966ca37ea", size = 71008, upload-time = "2025-10-12T14:55:18.883Z" },
 ]
 
-[[package]]
-name = "imageio-ffmpeg"
-version = "0.6.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/44/bd/c3343c721f2a1b0c9fc71c1aebf1966a3b7f08c2eea8ed5437a2865611d6/imageio_ffmpeg-0.6.0.tar.gz", hash = "sha256:e2556bed8e005564a9f925bb7afa4002d82770d6b08825078b7697ab88ba1755", size = 25210, upload-time = "2025-01-16T21:34:32.747Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/da/58/87ef68ac83f4c7690961bce288fd8e382bc5f1513860fc7f90a9c1c1c6bf/imageio_ffmpeg-0.6.0-py3-none-macosx_10_9_intel.macosx_10_9_x86_64.whl", hash = "sha256:9d2baaf867088508d4a3458e61eeb30e945c4ad8016025545f66c4b5aaef0a61", size = 24932969, upload-time = "2025-01-16T21:34:20.464Z" },
-    { url = "https://files.pythonhosted.org/packages/40/5c/f3d8a657d362cc93b81aab8feda487317da5b5d31c0e1fdfd5e986e55d17/imageio_ffmpeg-0.6.0-py3-none-macosx_11_0_arm64.whl", hash = "sha256:b1ae3173414b5fc5f538a726c4e48ea97edc0d2cdc11f103afee655c463fa742", size = 21113891, upload-time = "2025-01-16T21:34:00.277Z" },
-    { url = "https://files.pythonhosted.org/packages/33/e7/1925bfbc563c39c1d2e82501d8372734a5c725e53ac3b31b4c2d081e895b/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_aarch64.whl", hash = "sha256:1d47bebd83d2c5fc770720d211855f208af8a596c82d17730aa51e815cdee6dc", size = 25632706, upload-time = "2025-01-16T21:33:53.475Z" },
-    { url = "https://files.pythonhosted.org/packages/a0/2d/43c8522a2038e9d0e7dbdf3a61195ecc31ca576fb1527a528c877e87d973/imageio_ffmpeg-0.6.0-py3-none-manylinux2014_x86_64.whl", hash = "sha256:c7e46fcec401dd990405049d2e2f475e2b397779df2519b544b8aab515195282", size = 29498237, upload-time = "2025-01-16T21:34:13.726Z" },
-    { url = "https://files.pythonhosted.org/packages/a0/13/59da54728351883c3c1d9fca1710ab8eee82c7beba585df8f25ca925f08f/imageio_ffmpeg-0.6.0-py3-none-win32.whl", hash = "sha256:196faa79366b4a82f95c0f4053191d2013f4714a715780f0ad2a68ff37483cc2", size = 19652251, upload-time = "2025-01-16T21:34:06.812Z" },
-    { url = "https://files.pythonhosted.org/packages/2c/c6/fa760e12a2483469e2bf5058c5faff664acf66cadb4df2ad6205b016a73d/imageio_ffmpeg-0.6.0-py3-none-win_amd64.whl", hash = "sha256:02fa47c83703c37df6bfe4896aab339013f62bf02c5ebf2dce6da56af04ffc0a", size = 31246824, upload-time = "2025-01-16T21:34:28.6Z" },
-]
-
 [[package]]
 name = "importlib-metadata"
 version = "8.7.1"
@@ -781,6 +753,33 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl", hash = "sha256:5fc3c5039fc5ca8c0276333a188bbd59d6b7ab37fe6632daa76bc7f9ec18e713", size = 309071, upload-time = "2025-12-15T08:41:44.973Z" },
 ]
 
+[[package]]
+name = "jsonschema"
+version = "4.26.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "attrs" },
+    { name = "jsonschema-specifications" },
+    { name = "referencing" },
+    { name = "rpds-py" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/b3/fc/e067678238fa451312d4c62bf6e6cf5ec56375422aee02f9cb5f909b3047/jsonschema-4.26.0.tar.gz", hash = "sha256:0c26707e2efad8aa1bfc5b7ce170f3fccc2e4918ff85989ba9ffa9facb2be326", size = 366583, upload-time = "2026-01-07T13:41:07.246Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/69/90/f63fb5873511e014207a475e2bb4e8b2e570d655b00ac19a9a0ca0a385ee/jsonschema-4.26.0-py3-none-any.whl", hash = "sha256:d489f15263b8d200f8387e64b4c3a75f06629559fb73deb8fdfb525f2dab50ce", size = 90630, upload-time = "2026-01-07T13:41:05.306Z" },
+]
+
+[[package]]
+name = "jsonschema-specifications"
+version = "2025.9.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "referencing" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/19/74/a633ee74eb36c44aa6d1095e7cc5569bebf04342ee146178e2d36600708b/jsonschema_specifications-2025.9.1.tar.gz", hash = "sha256:b540987f239e745613c7a9176f3edb72b832a4ac465cf02712288397832b5e8d", size = 32855, upload-time = "2025-09-08T01:34:59.186Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/41/45/1a4ed80516f02155c51f51e8cedb3c1902296743db0bbc66608a0db2814f/jsonschema_specifications-2025.9.1-py3-none-any.whl", hash = "sha256:98802fee3a11ee76ecaca44429fda8a41bff98b00a0f2838151b113f210cc6fe", size = 18437, upload-time = "2025-09-08T01:34:57.871Z" },
+]
+
 [[package]]
 name = "kiwisolver"
 version = "1.5.0"
@@ -902,6 +901,40 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/91/63/cd3ca7a16fc0201069ea951615431dbeae489e348f2f4137cff957fb750d/megatron_fsdp-0.3.1-py3-none-any.whl", hash = "sha256:290b95e530f530a953909dc73b9e0fa480b3c660f2538450e83bbaa0215af247", size = 94391, upload-time = "2026-03-20T21:23:35.75Z" },
 ]
 
+[[package]]
+name = "mistral-common"
+version = "1.11.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "jsonschema" },
+    { name = "numpy" },
+    { name = "pillow" },
+    { name = "pydantic" },
+    { name = "pydantic-extra-types", extra = ["pycountry"] },
+    { name = "requests" },
+    { name = "tiktoken" },
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/c2/eb/12167a1bea9714582e5b4f539f9c019323363e314a499c72855ff0e5ad43/mistral_common-1.11.2.tar.gz", hash = "sha256:79f68fc2d1190f28637f40e053f919c8c2697e00b2aa679ddee562a95183f4ad", size = 6357845, upload-time = "2026-05-04T19:47:40.413Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/47/f0/6a5d604b972e442b9d36c117d01788feddad099e4965699e3516ee6fefc3/mistral_common-1.11.2-py3-none-any.whl", hash = "sha256:ebb42062cd705a0aa2bc69b4cde2b83d446ae58150b7e29322c90cb08fcfca6c", size = 6531968, upload-time = "2026-05-04T19:47:37.718Z" },
+]
+
+[package.optional-dependencies]
+audio = [
+    { name = "soundfile" },
+    { name = "soxr" },
+]
+hf-hub = [
+    { name = "huggingface-hub" },
+]
+image = [
+    { name = "opencv-python-headless" },
+]
+sentencepiece = [
+    { name = "sentencepiece" },
+]
+
 [[package]]
 name = "mlflow"
 version = "3.10.1"
@@ -1034,18 +1067,18 @@ wheels = [
 
 [[package]]
 name = "nemo-automodel"
-version = "0.3.0rc0"
-source = { git = "https://github.com/NVIDIA-NeMo/Automodel.git?rev=ecd7cb4297ae06990c94b7581f34c03ce8ed5488#ecd7cb4297ae06990c94b7581f34c03ce8ed5488" }
+version = "0.4.0"
+source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "datasets" },
-    { name = "diffusers" },
-    { name = "ftfy" },
-    { name = "imageio-ffmpeg" },
+    { name = "flashoptim" },
     { name = "megatron-fsdp" },
+    { name = "mistral-common", extra = ["audio", "hf-hub", "image", "sentencepiece"] },
     { name = "mlflow" },
     { name = "opencv-python-headless" },
     { name = "pybind11" },
     { name = "pyyaml" },
+    { name = "tiktoken" },
     { name = "torch", version = "2.10.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin'" },
     { name = "torch", version = "2.10.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform != 'darwin' and sys_platform != 'linux'" },
     { name = "torch", version = "2.10.0+cu129", source = { registry = "https://download.pytorch.org/whl/cu129" }, marker = "sys_platform == 'linux'" },
@@ -1054,6 +1087,10 @@ dependencies = [
     { name = "transformers" },
     { name = "wandb" },
 ]
+sdist = { url = "https://files.pythonhosted.org/packages/f7/2b/303821a87b340dc0466517e1728636d6296cfc7a00ac4dcc7406763cd38b/nemo_automodel-0.4.0.tar.gz", hash = "sha256:407dcdf1251ea608f3cb1f7d4a9b79b9401cfbc53386172486d69535aa9d8cb2", size = 879465, upload-time = "2026-04-28T20:04:14.117Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/27/7e/bec1f6444c63149004e50f5912d0718a6d11c1922bec2617d50e10788e51/nemo_automodel-0.4.0-py3-none-any.whl", hash = "sha256:08e66717eabb4288eba437f1e53e85abc6f0f7f0b19264f09dcb338c9d5487a2", size = 1106299, upload-time = "2026-04-28T20:04:12.35Z" },
+]
 
 [[package]]
 name = "networkx"
@@ -1066,21 +1103,21 @@ wheels = [
 
 [[package]]
 name = "numpy"
-version = "2.4.4"
+version = "2.3.5"
 source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/d7/9f/b8cef5bffa569759033adda9481211426f12f53299629b410340795c2514/numpy-2.4.4.tar.gz", hash = "sha256:2d390634c5182175533585cc89f3608a4682ccb173cc9bb940b2881c8d6f8fa0", size = 20731587, upload-time = "2026-03-29T13:22:01.298Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/76/65/21b3bc86aac7b8f2862db1e808f1ea22b028e30a225a34a5ede9bf8678f2/numpy-2.3.5.tar.gz", hash = "sha256:784db1dcdab56bf0517743e746dfb0f885fc68d948aba86eeec2cba234bdf1c0", size = 20584950, upload-time = "2025-11-16T22:52:42.067Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/28/05/32396bec30fb2263770ee910142f49c1476d08e8ad41abf8403806b520ce/numpy-2.4.4-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:15716cfef24d3a9762e3acdf87e27f58dc823d1348f765bbea6bef8c639bfa1b", size = 16689272, upload-time = "2026-03-29T13:18:49.223Z" },
-    { url = "https://files.pythonhosted.org/packages/c5/f3/a983d28637bfcd763a9c7aafdb6d5c0ebf3d487d1e1459ffdb57e2f01117/numpy-2.4.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:23cbfd4c17357c81021f21540da84ee282b9c8fba38a03b7b9d09ba6b951421e", size = 14699573, upload-time = "2026-03-29T13:18:52.629Z" },
-    { url = "https://files.pythonhosted.org/packages/9b/fd/e5ecca1e78c05106d98028114f5c00d3eddb41207686b2b7de3e477b0e22/numpy-2.4.4-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:8b3b60bb7cba2c8c81837661c488637eee696f59a877788a396d33150c35d842", size = 5204782, upload-time = "2026-03-29T13:18:55.579Z" },
-    { url = "https://files.pythonhosted.org/packages/de/2f/702a4594413c1a8632092beae8aba00f1d67947389369b3777aed783fdca/numpy-2.4.4-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:e4a010c27ff6f210ff4c6ef34394cd61470d01014439b192ec22552ee867f2a8", size = 6552038, upload-time = "2026-03-29T13:18:57.769Z" },
-    { url = "https://files.pythonhosted.org/packages/7f/37/eed308a8f56cba4d1fdf467a4fc67ef4ff4bf1c888f5fc980481890104b1/numpy-2.4.4-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:f9e75681b59ddaa5e659898085ae0eaea229d054f2ac0c7e563a62205a700121", size = 15670666, upload-time = "2026-03-29T13:19:00.341Z" },
-    { url = "https://files.pythonhosted.org/packages/0a/0d/0e3ecece05b7a7e87ab9fb587855548da437a061326fff64a223b6dcb78a/numpy-2.4.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:81f4a14bee47aec54f883e0cad2d73986640c1590eb9bfaaba7ad17394481e6e", size = 16645480, upload-time = "2026-03-29T13:19:03.63Z" },
-    { url = "https://files.pythonhosted.org/packages/34/49/f2312c154b82a286758ee2f1743336d50651f8b5195db18cdb63675ff649/numpy-2.4.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:62d6b0f03b694173f9fcb1fb317f7222fd0b0b103e784c6549f5e53a27718c44", size = 17020036, upload-time = "2026-03-29T13:19:07.428Z" },
-    { url = "https://files.pythonhosted.org/packages/7b/e9/736d17bd77f1b0ec4f9901aaec129c00d59f5d84d5e79bba540ef12c2330/numpy-2.4.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:fbc356aae7adf9e6336d336b9c8111d390a05df88f1805573ebb0807bd06fd1d", size = 18368643, upload-time = "2026-03-29T13:19:10.775Z" },
-    { url = "https://files.pythonhosted.org/packages/63/f6/d417977c5f519b17c8a5c3bc9e8304b0908b0e21136fe43bf628a1343914/numpy-2.4.4-cp312-cp312-win32.whl", hash = "sha256:0d35aea54ad1d420c812bfa0385c71cd7cc5bcf7c65fed95fc2cd02fe8c79827", size = 5961117, upload-time = "2026-03-29T13:19:13.464Z" },
-    { url = "https://files.pythonhosted.org/packages/2d/5b/e1deebf88ff431b01b7406ca3583ab2bbb90972bbe1c568732e49c844f7e/numpy-2.4.4-cp312-cp312-win_amd64.whl", hash = "sha256:b5f0362dc928a6ecd9db58868fca5e48485205e3855957bdedea308f8672ea4a", size = 12320584, upload-time = "2026-03-29T13:19:16.155Z" },
-    { url = "https://files.pythonhosted.org/packages/58/89/e4e856ac82a68c3ed64486a544977d0e7bdd18b8da75b78a577ca31c4395/numpy-2.4.4-cp312-cp312-win_arm64.whl", hash = "sha256:846300f379b5b12cc769334464656bc882e0735d27d9726568bc932fdc49d5ec", size = 10221450, upload-time = "2026-03-29T13:19:18.994Z" },
+    { url = "https://files.pythonhosted.org/packages/44/37/e669fe6cbb2b96c62f6bbedc6a81c0f3b7362f6a59230b23caa673a85721/numpy-2.3.5-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:74ae7b798248fe62021dbf3c914245ad45d1a6b0cb4a29ecb4b31d0bfbc4cc3e", size = 16733873, upload-time = "2025-11-16T22:49:49.84Z" },
+    { url = "https://files.pythonhosted.org/packages/c5/65/df0db6c097892c9380851ab9e44b52d4f7ba576b833996e0080181c0c439/numpy-2.3.5-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:ee3888d9ff7c14604052b2ca5535a30216aa0a58e948cdd3eeb8d3415f638769", size = 12259838, upload-time = "2025-11-16T22:49:52.863Z" },
+    { url = "https://files.pythonhosted.org/packages/5b/e1/1ee06e70eb2136797abe847d386e7c0e830b67ad1d43f364dd04fa50d338/numpy-2.3.5-cp312-cp312-macosx_14_0_arm64.whl", hash = "sha256:612a95a17655e213502f60cfb9bf9408efdc9eb1d5f50535cc6eb365d11b42b5", size = 5088378, upload-time = "2025-11-16T22:49:55.055Z" },
+    { url = "https://files.pythonhosted.org/packages/6d/9c/1ca85fb86708724275103b81ec4cf1ac1d08f465368acfc8da7ab545bdae/numpy-2.3.5-cp312-cp312-macosx_14_0_x86_64.whl", hash = "sha256:3101e5177d114a593d79dd79658650fe28b5a0d8abeb8ce6f437c0e6df5be1a4", size = 6628559, upload-time = "2025-11-16T22:49:57.371Z" },
+    { url = "https://files.pythonhosted.org/packages/74/78/fcd41e5a0ce4f3f7b003da85825acddae6d7ecb60cf25194741b036ca7d6/numpy-2.3.5-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8b973c57ff8e184109db042c842423ff4f60446239bd585a5131cc47f06f789d", size = 14250702, upload-time = "2025-11-16T22:49:59.632Z" },
+    { url = "https://files.pythonhosted.org/packages/b6/23/2a1b231b8ff672b4c450dac27164a8b2ca7d9b7144f9c02d2396518352eb/numpy-2.3.5-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0d8163f43acde9a73c2a33605353a4f1bc4798745a8b1d73183b28e5b435ae28", size = 16606086, upload-time = "2025-11-16T22:50:02.127Z" },
+    { url = "https://files.pythonhosted.org/packages/a0/c5/5ad26fbfbe2012e190cc7d5003e4d874b88bb18861d0829edc140a713021/numpy-2.3.5-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:51c1e14eb1e154ebd80e860722f9e6ed6ec89714ad2db2d3aa33c31d7c12179b", size = 16025985, upload-time = "2025-11-16T22:50:04.536Z" },
+    { url = "https://files.pythonhosted.org/packages/d2/fa/dd48e225c46c819288148d9d060b047fd2a6fb1eb37eae25112ee4cb4453/numpy-2.3.5-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:b46b4ec24f7293f23adcd2d146960559aaf8020213de8ad1909dba6c013bf89c", size = 18542976, upload-time = "2025-11-16T22:50:07.557Z" },
+    { url = "https://files.pythonhosted.org/packages/05/79/ccbd23a75862d95af03d28b5c6901a1b7da4803181513d52f3b86ed9446e/numpy-2.3.5-cp312-cp312-win32.whl", hash = "sha256:3997b5b3c9a771e157f9aae01dd579ee35ad7109be18db0e85dbdbe1de06e952", size = 6285274, upload-time = "2025-11-16T22:50:10.746Z" },
+    { url = "https://files.pythonhosted.org/packages/2d/57/8aeaf160312f7f489dea47ab61e430b5cb051f59a98ae68b7133ce8fa06a/numpy-2.3.5-cp312-cp312-win_amd64.whl", hash = "sha256:86945f2ee6d10cdfd67bcb4069c1662dd711f7e2a4343db5cecec06b87cf31aa", size = 12782922, upload-time = "2025-11-16T22:50:12.811Z" },
+    { url = "https://files.pythonhosted.org/packages/78/a6/aae5cc2ca78c45e64b9ef22f089141d661516856cf7c8a54ba434576900d/numpy-2.3.5-cp312-cp312-win_arm64.whl", hash = "sha256:f28620fe26bee16243be2b7b874da327312240a7cdc38b769a697578d2100013", size = 10194667, upload-time = "2025-11-16T22:50:16.16Z" },
 ]
 
 [[package]]
@@ -1468,6 +1505,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/88/c5/e98d9c51f3d5300d5e40ad9037dd6b3b60736fd02ab68dcc98c96be7592d/pybind11-3.0.2-py3-none-any.whl", hash = "sha256:f8a6500548919cc33bcd220d5f984688326f574fa97f1107f2f4fdb4c6fb019f", size = 310158, upload-time = "2026-02-17T04:46:49.91Z" },
 ]
 
+[[package]]
+name = "pycountry"
+version = "26.2.16"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/de/1d/061b9e7a48b85cfd69f33c33d2ef784a531c359399ad764243399673c8f5/pycountry-26.2.16.tar.gz", hash = "sha256:5b6027d453fcd6060112b951dd010f01f168b51b4bf8a1f1fc8c95c8d94a0801", size = 7711342, upload-time = "2026-02-17T03:42:52.367Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/9c/42/7703bd45b62fecd44cd7d3495423097e2f7d28bc2e99e7c1af68892ab157/pycountry-26.2.16-py3-none-any.whl", hash = "sha256:115c4baf7cceaa30f59a4694d79483c9167dbce7a9de4d3d571c5f3ea77c305a", size = 8044600, upload-time = "2026-02-17T03:42:49.777Z" },
+]
+
 [[package]]
 name = "pycparser"
 version = "3.0"
@@ -1521,6 +1567,24 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/f7/07/34573da085946b6a313d7c42f82f16e8920bfd730665de2d11c0c37a74b5/pydantic_core-2.41.5-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:76d0819de158cd855d1cbb8fcafdf6f5cf1eb8e470abe056d5d161106e38062b", size = 2139017, upload-time = "2025-11-04T13:42:59.471Z" },
 ]
 
+[[package]]
+name = "pydantic-extra-types"
+version = "2.11.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "pydantic" },
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/66/71/dba38ee2651f84f7842206adbd2233d8bbdb59fb85e9fa14232486a8c471/pydantic_extra_types-2.11.1.tar.gz", hash = "sha256:46792d2307383859e923d8fcefa82108b1a141f8a9c0198982b3832ab5ef1049", size = 172002, upload-time = "2026-03-16T08:08:03.92Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/17/c1/3226e6d7f5a4f736f38ac11a6fbb262d701889802595cdb0f53a885ac2e0/pydantic_extra_types-2.11.1-py3-none-any.whl", hash = "sha256:1722ea2bddae5628ace25f2aa685b69978ef533123e5638cfbddb999e0100ec1", size = 79526, upload-time = "2026-03-16T08:08:02.533Z" },
+]
+
+[package.optional-dependencies]
+pycountry = [
+    { name = "pycountry" },
+]
+
 [[package]]
 name = "pydantic-settings"
 version = "2.13.1"
@@ -1622,18 +1686,30 @@ dependencies = [
     { name = "torch", version = "2.10.0", source = { registry = "https://pypi.org/simple" }, marker = "sys_platform == 'darwin'" },
     { name = "torch", version = "2.10.0+cpu", source = { registry = "https://download.pytorch.org/whl/cpu" }, marker = "sys_platform != 'darwin' and sys_platform != 'linux'" },
     { name = "torch", version = "2.10.0+cu129", source = { registry = "https://download.pytorch.org/whl/cu129" }, marker = "sys_platform == 'linux'" },
-    { name = "transformers" },
 ]
 
 [package.metadata]
 requires-dist = [
-    { name = "nemo-automodel", git = "https://github.com/NVIDIA-NeMo/Automodel.git?rev=ecd7cb4297ae06990c94b7581f34c03ce8ed5488" },
+    { name = "nemo-automodel", specifier = "==0.4.0" },
     { name = "omegaconf", specifier = ">=2.3.0" },
     { name = "pydantic-settings", specifier = ">=2.0.0" },
     { name = "torch", marker = "sys_platform != 'darwin' and sys_platform != 'linux'", specifier = ">=2.10.0", index = "https://download.pytorch.org/whl/cpu" },
     { name = "torch", marker = "sys_platform == 'darwin'", specifier = ">=2.10.0" },
     { name = "torch", marker = "sys_platform == 'linux'", specifier = ">=2.10.0", index = "https://download.pytorch.org/whl/cu129" },
-    { name = "transformers", specifier = ">=5.0,<5.2" },
+]
+
+[[package]]
+name = "referencing"
+version = "0.37.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "attrs" },
+    { name = "rpds-py" },
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/22/f5/df4e9027acead3ecc63e50fe1e36aca1523e1719559c499951bb4b53188f/referencing-0.37.0.tar.gz", hash = "sha256:44aefc3142c5b842538163acb373e24cce6632bd54bdb01b21ad5863489f50d8", size = 78036, upload-time = "2025-10-13T15:30:48.871Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/2c/58/ca301544e1fa93ed4f80d724bf5b194f6e4b945841c5bfd555878eea9fcb/referencing-0.37.0-py3-none-any.whl", hash = "sha256:381329a9f99628c9069361716891d34ad94af76e461dcb0335825aecc7692231", size = 26766, upload-time = "2025-10-13T15:30:47.625Z" },
 ]
 
 [[package]]
@@ -1688,6 +1764,29 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/14/25/b208c5683343959b670dc001595f2f3737e051da617f66c31f7c4fa93abc/rich-14.3.3-py3-none-any.whl", hash = "sha256:793431c1f8619afa7d3b52b2cdec859562b950ea0d4b6b505397612db8d5362d", size = 310458, upload-time = "2026-02-19T17:23:13.732Z" },
 ]
 
+[[package]]
+name = "rpds-py"
+version = "0.30.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/20/af/3f2f423103f1113b36230496629986e0ef7e199d2aa8392452b484b38ced/rpds_py-0.30.0.tar.gz", hash = "sha256:dd8ff7cf90014af0c0f787eea34794ebf6415242ee1d6fa91eaba725cc441e84", size = 69469, upload-time = "2025-11-30T20:24:38.837Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/03/e7/98a2f4ac921d82f33e03f3835f5bf3a4a40aa1bfdc57975e74a97b2b4bdd/rpds_py-0.30.0-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:a161f20d9a43006833cd7068375a94d035714d73a172b681d8881820600abfad", size = 375086, upload-time = "2025-11-30T20:22:17.93Z" },
+    { url = "https://files.pythonhosted.org/packages/4d/a1/bca7fd3d452b272e13335db8d6b0b3ecde0f90ad6f16f3328c6fb150c889/rpds_py-0.30.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:6abc8880d9d036ecaafe709079969f56e876fcf107f7a8e9920ba6d5a3878d05", size = 359053, upload-time = "2025-11-30T20:22:19.297Z" },
+    { url = "https://files.pythonhosted.org/packages/65/1c/ae157e83a6357eceff62ba7e52113e3ec4834a84cfe07fa4b0757a7d105f/rpds_py-0.30.0-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ca28829ae5f5d569bb62a79512c842a03a12576375d5ece7d2cadf8abe96ec28", size = 390763, upload-time = "2025-11-30T20:22:21.661Z" },
+    { url = "https://files.pythonhosted.org/packages/d4/36/eb2eb8515e2ad24c0bd43c3ee9cd74c33f7ca6430755ccdb240fd3144c44/rpds_py-0.30.0-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:a1010ed9524c73b94d15919ca4d41d8780980e1765babf85f9a2f90d247153dd", size = 408951, upload-time = "2025-11-30T20:22:23.408Z" },
+    { url = "https://files.pythonhosted.org/packages/d6/65/ad8dc1784a331fabbd740ef6f71ce2198c7ed0890dab595adb9ea2d775a1/rpds_py-0.30.0-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:f8d1736cfb49381ba528cd5baa46f82fdc65c06e843dab24dd70b63d09121b3f", size = 514622, upload-time = "2025-11-30T20:22:25.16Z" },
+    { url = "https://files.pythonhosted.org/packages/63/8e/0cfa7ae158e15e143fe03993b5bcd743a59f541f5952e1546b1ac1b5fd45/rpds_py-0.30.0-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:d948b135c4693daff7bc2dcfc4ec57237a29bd37e60c2fabf5aff2bbacf3e2f1", size = 414492, upload-time = "2025-11-30T20:22:26.505Z" },
+    { url = "https://files.pythonhosted.org/packages/60/1b/6f8f29f3f995c7ffdde46a626ddccd7c63aefc0efae881dc13b6e5d5bb16/rpds_py-0.30.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:47f236970bccb2233267d89173d3ad2703cd36a0e2a6e92d0560d333871a3d23", size = 394080, upload-time = "2025-11-30T20:22:27.934Z" },
+    { url = "https://files.pythonhosted.org/packages/6d/d5/a266341051a7a3ca2f4b750a3aa4abc986378431fc2da508c5034d081b70/rpds_py-0.30.0-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:2e6ecb5a5bcacf59c3f912155044479af1d0b6681280048b338b28e364aca1f6", size = 408680, upload-time = "2025-11-30T20:22:29.341Z" },
+    { url = "https://files.pythonhosted.org/packages/10/3b/71b725851df9ab7a7a4e33cf36d241933da66040d195a84781f49c50490c/rpds_py-0.30.0-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:a8fa71a2e078c527c3e9dc9fc5a98c9db40bcc8a92b4e8858e36d329f8684b51", size = 423589, upload-time = "2025-11-30T20:22:31.469Z" },
+    { url = "https://files.pythonhosted.org/packages/00/2b/e59e58c544dc9bd8bd8384ecdb8ea91f6727f0e37a7131baeff8d6f51661/rpds_py-0.30.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:73c67f2db7bc334e518d097c6d1e6fed021bbc9b7d678d6cc433478365d1d5f5", size = 573289, upload-time = "2025-11-30T20:22:32.997Z" },
+    { url = "https://files.pythonhosted.org/packages/da/3e/a18e6f5b460893172a7d6a680e86d3b6bc87a54c1f0b03446a3c8c7b588f/rpds_py-0.30.0-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:5ba103fb455be00f3b1c2076c9d4264bfcb037c976167a6047ed82f23153f02e", size = 599737, upload-time = "2025-11-30T20:22:34.419Z" },
+    { url = "https://files.pythonhosted.org/packages/5c/e2/714694e4b87b85a18e2c243614974413c60aa107fd815b8cbc42b873d1d7/rpds_py-0.30.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:7cee9c752c0364588353e627da8a7e808a66873672bcb5f52890c33fd965b394", size = 563120, upload-time = "2025-11-30T20:22:35.903Z" },
+    { url = "https://files.pythonhosted.org/packages/6f/ab/d5d5e3bcedb0a77f4f613706b750e50a5a3ba1c15ccd3665ecc636c968fd/rpds_py-0.30.0-cp312-cp312-win32.whl", hash = "sha256:1ab5b83dbcf55acc8b08fc62b796ef672c457b17dbd7820a11d6c52c06839bdf", size = 223782, upload-time = "2025-11-30T20:22:37.271Z" },
+    { url = "https://files.pythonhosted.org/packages/39/3b/f786af9957306fdc38a74cef405b7b93180f481fb48453a114bb6465744a/rpds_py-0.30.0-cp312-cp312-win_amd64.whl", hash = "sha256:a090322ca841abd453d43456ac34db46e8b05fd9b3b4ac0c78bcde8b089f959b", size = 240463, upload-time = "2025-11-30T20:22:39.021Z" },
+    { url = "https://files.pythonhosted.org/packages/f3/d2/b91dc748126c1559042cfe41990deb92c4ee3e2b415f6b5234969ffaf0cc/rpds_py-0.30.0-cp312-cp312-win_arm64.whl", hash = "sha256:669b1805bd639dd2989b281be2cfd951c6121b65e729d9b843e9639ef1fd555e", size = 230868, upload-time = "2025-11-30T20:22:40.493Z" },
+]
+
 [[package]]
 name = "safetensors"
 version = "0.7.0"
@@ -1751,6 +1850,22 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/bc/98/fe9ae9ffb3b54b62559f52dedaebe204b408db8109a8c66fdd04869e6424/scipy-1.17.1-cp312-cp312-win_arm64.whl", hash = "sha256:f4115102802df98b2b0db3cce5cb9b92572633a1197c77b7553e5203f284a5b3", size = 24547340, upload-time = "2026-02-23T00:19:12.024Z" },
 ]
 
+[[package]]
+name = "sentencepiece"
+version = "0.2.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/15/15/2e7a025fc62d764b151ae6d0f2a92f8081755ebe8d4a64099accc6f77ba6/sentencepiece-0.2.1.tar.gz", hash = "sha256:8138cec27c2f2282f4a34d9a016e3374cd40e5c6e9cb335063db66a0a3b71fad", size = 3228515, upload-time = "2025-08-12T07:00:51.718Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/4a/be/32ce495aa1d0e0c323dcb1ba87096037358edee539cac5baf8755a6bd396/sentencepiece-0.2.1-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:57cae326c8727de58c85977b175af132a7138d84c764635d7e71bbee7e774133", size = 1943152, upload-time = "2025-08-12T06:59:40.048Z" },
+    { url = "https://files.pythonhosted.org/packages/88/7e/ff23008899a58678e98c6ff592bf4d368eee5a71af96d0df6b38a039dd4f/sentencepiece-0.2.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:56dd39a3c4d6493db3cdca7e8cc68c6b633f0d4195495cbadfcf5af8a22d05a6", size = 1325651, upload-time = "2025-08-12T06:59:41.536Z" },
+    { url = "https://files.pythonhosted.org/packages/19/84/42eb3ce4796777a1b5d3699dfd4dca85113e68b637f194a6c8d786f16a04/sentencepiece-0.2.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d9381351182ff9888cc80e41c632e7e274b106f450de33d67a9e8f6043da6f76", size = 1253645, upload-time = "2025-08-12T06:59:42.903Z" },
+    { url = "https://files.pythonhosted.org/packages/89/fa/d3d5ebcba3cb9e6d3775a096251860c41a6bc53a1b9461151df83fe93255/sentencepiece-0.2.1-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:99f955df238021bf11f0fc37cdb54fd5e5b5f7fd30ecc3d93fb48b6815437167", size = 1316273, upload-time = "2025-08-12T06:59:44.476Z" },
+    { url = "https://files.pythonhosted.org/packages/04/88/14f2f4a2b922d8b39be45bf63d79e6cd3a9b2f248b2fcb98a69b12af12f5/sentencepiece-0.2.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0cdfecef430d985f1c2bcbfff3defd1d95dae876fbd0173376012d2d7d24044b", size = 1387881, upload-time = "2025-08-12T06:59:46.09Z" },
+    { url = "https://files.pythonhosted.org/packages/fd/b8/903e5ccb77b4ef140605d5d71b4f9e0ad95d456d6184688073ed11712809/sentencepiece-0.2.1-cp312-cp312-win32.whl", hash = "sha256:a483fd29a34c3e34c39ac5556b0a90942bec253d260235729e50976f5dba1068", size = 999540, upload-time = "2025-08-12T06:59:48.023Z" },
+    { url = "https://files.pythonhosted.org/packages/2d/81/92df5673c067148c2545b1bfe49adfd775bcc3a169a047f5a0e6575ddaca/sentencepiece-0.2.1-cp312-cp312-win_amd64.whl", hash = "sha256:4cdc7c36234fda305e85c32949c5211faaf8dd886096c7cea289ddc12a2d02de", size = 1054671, upload-time = "2025-08-12T06:59:49.895Z" },
+    { url = "https://files.pythonhosted.org/packages/fe/02/c5e3bc518655d714622bec87d83db9cdba1cd0619a4a04e2109751c4f47f/sentencepiece-0.2.1-cp312-cp312-win_arm64.whl", hash = "sha256:daeb5e9e9fcad012324807856113708614d534f596d5008638eb9b40112cd9e4", size = 1033923, upload-time = "2025-08-12T06:59:51.952Z" },
+]
+
 [[package]]
 name = "sentry-sdk"
 version = "2.56.0"
@@ -1816,6 +1931,41 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/c1/d4/59e74daffcb57a07668852eeeb6035af9f32cbfd7a1d2511f17d2fe6a738/smmap-5.0.3-py3-none-any.whl", hash = "sha256:c106e05d5a61449cf6ba9a1e650227ecfb141590d2a98412103ff35d89fc7b2f", size = 24390, upload-time = "2026-03-09T03:43:24.361Z" },
 ]
 
+[[package]]
+name = "soundfile"
+version = "0.13.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "cffi" },
+    { name = "numpy" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/e1/41/9b873a8c055582859b239be17902a85339bec6a30ad162f98c9b0288a2cc/soundfile-0.13.1.tar.gz", hash = "sha256:b2c68dab1e30297317080a5b43df57e302584c49e2942defdde0acccc53f0e5b", size = 46156, upload-time = "2025-01-25T09:17:04.831Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/64/28/e2a36573ccbcf3d57c00626a21fe51989380636e821b341d36ccca0c1c3a/soundfile-0.13.1-py2.py3-none-any.whl", hash = "sha256:a23c717560da2cf4c7b5ae1142514e0fd82d6bbd9dfc93a50423447142f2c445", size = 25751, upload-time = "2025-01-25T09:16:44.235Z" },
+    { url = "https://files.pythonhosted.org/packages/ea/ab/73e97a5b3cc46bba7ff8650a1504348fa1863a6f9d57d7001c6b67c5f20e/soundfile-0.13.1-py2.py3-none-macosx_10_9_x86_64.whl", hash = "sha256:82dc664d19831933fe59adad199bf3945ad06d84bc111a5b4c0d3089a5b9ec33", size = 1142250, upload-time = "2025-01-25T09:16:47.583Z" },
+    { url = "https://files.pythonhosted.org/packages/a0/e5/58fd1a8d7b26fc113af244f966ee3aecf03cb9293cb935daaddc1e455e18/soundfile-0.13.1-py2.py3-none-macosx_11_0_arm64.whl", hash = "sha256:743f12c12c4054921e15736c6be09ac26b3b3d603aef6fd69f9dde68748f2593", size = 1101406, upload-time = "2025-01-25T09:16:49.662Z" },
+    { url = "https://files.pythonhosted.org/packages/58/ae/c0e4a53d77cf6e9a04179535766b3321b0b9ced5f70522e4caf9329f0046/soundfile-0.13.1-py2.py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:9c9e855f5a4d06ce4213f31918653ab7de0c5a8d8107cd2427e44b42df547deb", size = 1235729, upload-time = "2025-01-25T09:16:53.018Z" },
+    { url = "https://files.pythonhosted.org/packages/57/5e/70bdd9579b35003a489fc850b5047beeda26328053ebadc1fb60f320f7db/soundfile-0.13.1-py2.py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:03267c4e493315294834a0870f31dbb3b28a95561b80b134f0bd3cf2d5f0e618", size = 1313646, upload-time = "2025-01-25T09:16:54.872Z" },
+    { url = "https://files.pythonhosted.org/packages/fe/df/8c11dc4dfceda14e3003bb81a0d0edcaaf0796dd7b4f826ea3e532146bba/soundfile-0.13.1-py2.py3-none-win32.whl", hash = "sha256:c734564fab7c5ddf8e9be5bf70bab68042cd17e9c214c06e365e20d64f9a69d5", size = 899881, upload-time = "2025-01-25T09:16:56.663Z" },
+    { url = "https://files.pythonhosted.org/packages/14/e9/6b761de83277f2f02ded7e7ea6f07828ec78e4b229b80e4ca55dd205b9dc/soundfile-0.13.1-py2.py3-none-win_amd64.whl", hash = "sha256:1e70a05a0626524a69e9f0f4dd2ec174b4e9567f4d8b6c11d38b5c289be36ee9", size = 1019162, upload-time = "2025-01-25T09:16:59.573Z" },
+]
+
+[[package]]
+name = "soxr"
+version = "1.1.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "numpy" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/ed/11/27cebce4a108f77afea7c80545115536b45e3f11ebfb914f638fdd9ba847/soxr-1.1.0.tar.gz", hash = "sha256:9f228ae21c78fa9359ca98d8a5e8e91f30639e438e574133dace62c5b5309e44", size = 173067, upload-time = "2026-05-03T00:15:18.214Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/06/8a/f3da7973b5f1b05d2d7e94d5376b881dcbc05297900cae6c3d33d95b209b/soxr-1.1.0-cp312-abi3-macosx_10_14_x86_64.whl", hash = "sha256:e0e09fa633ce2e67df08b298afced4d184f6e753fc330f241022250f1d0d61da", size = 204124, upload-time = "2026-05-03T00:14:54.505Z" },
+    { url = "https://files.pythonhosted.org/packages/03/dc/200013a74641f8774664bbcd2346c695c05c2e300ea792adcb40a293eed0/soxr-1.1.0-cp312-abi3-macosx_11_0_arm64.whl", hash = "sha256:d6a7ad82b8d5f3fcc04b1d2ca055562b96af571e1d4fa7c6c61d0fb509ac43b4", size = 165457, upload-time = "2026-05-03T00:14:56.007Z" },
+    { url = "https://files.pythonhosted.org/packages/88/2b/2e5eba817a762a2ec589ff165b8bc5955b25a0ad140045f7cd8e45410543/soxr-1.1.0-cp312-abi3-manylinux_2_26_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bf98c0d7b7d5ef5bf072fee8d3020e8b664f2d195933ea7bc5089267c2e22a06", size = 206529, upload-time = "2026-05-03T00:14:57.646Z" },
+    { url = "https://files.pythonhosted.org/packages/5c/f1/0e55195893228609c9a08c3b13b7a83a46c3a992cd00d3304f0f320cfb07/soxr-1.1.0-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3b033078e86f3c4a658e5697fac8995764fad9e799563616b630136b613167f1", size = 240413, upload-time = "2026-05-03T00:14:59.363Z" },
+    { url = "https://files.pythonhosted.org/packages/b0/4d/621e4150e4815246ad552d215a8a294a90143fedd19ee442cf82d3b3abc8/soxr-1.1.0-cp312-abi3-win_amd64.whl", hash = "sha256:6ae2a174bffea94e8ead857dad85999d3f49f091774dbad5b046c0417d7092f4", size = 174357, upload-time = "2026-05-03T00:15:00.724Z" },
+]
+
 [[package]]
 name = "sqlalchemy"
 version = "2.0.48"
@@ -1879,6 +2029,25 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/32/d5/f9a850d79b0851d1d4ef6456097579a9005b31fea68726a4ae5f2d82ddd9/threadpoolctl-3.6.0-py3-none-any.whl", hash = "sha256:43a0b8fd5a2928500110039e43a5eed8480b918967083ea48dc3ab9f13c4a7fb", size = 18638, upload-time = "2025-03-13T13:49:21.846Z" },
 ]
 
+[[package]]
+name = "tiktoken"
+version = "0.13.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "regex" },
+    { name = "requests" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/e4/e5/5f3cb2159769d0f4324c0e9e87f9de3c4b1cd45848a96b2eb3566ad5ca77/tiktoken-0.13.0.tar.gz", hash = "sha256:c9435714c3a84c2319499de9a300c0e604449dd0799ff246458b3bb6a7f433c1", size = 38986, upload-time = "2026-05-15T04:51:27.153Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/85/8e/144bde4e01df66b34bb865557c7cd754ed08b036217ebd79c9db5e9048a9/tiktoken-0.13.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:32ac870a806cfb260a02d0cb70426aef02e038297f8ad50df5040bb5af360791", size = 1034888, upload-time = "2026-05-15T04:50:31.579Z" },
+    { url = "https://files.pythonhosted.org/packages/36/18/d4ac9d20956cdebca04841316660ed584c2fecdc2b81722a28bc7ad3b1e4/tiktoken-0.13.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:4d9980f11429ed2d737c463bb1fb78cf330caa026adf002f714aced7849a687b", size = 982970, upload-time = "2026-05-15T04:50:32.961Z" },
+    { url = "https://files.pythonhosted.org/packages/74/ed/6bb8d05b9f731f749fee5c6f5ca63e981143c826a5985877330507bd13b7/tiktoken-0.13.0-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:3f277ebea5edd7b8bf03c6f9431e1d67d517530115572b2dc1d465326e8f88c7", size = 1115741, upload-time = "2026-05-15T04:50:34.475Z" },
+    { url = "https://files.pythonhosted.org/packages/34/de/2ca96b07a82d972b74fe4b46de055b79c904e45c7eab699354a0bfa697dc/tiktoken-0.13.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:a116178fa7e1b4065bff05214360373a65cac22f965be7b3f73d00a0dbfe7649", size = 1136523, upload-time = "2026-05-15T04:50:35.782Z" },
+    { url = "https://files.pythonhosted.org/packages/ee/dc/9dafec002c2d4424378563cf4cf5c7fb93631d2a55013c8b87554ee4012c/tiktoken-0.13.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:2c397ddda233208345b01bd30f2fca79ff730e55731d0108a603f9bc57f6af3b", size = 1181954, upload-time = "2026-05-15T04:50:36.99Z" },
+    { url = "https://files.pythonhosted.org/packages/a1/d0/1f8578c45b2f24759b46f0b50d31878c63c73e6bf0f2227e10ec5c5408dc/tiktoken-0.13.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:95097e4f89b06403976e498abf61a0ee73a7497e73fb599cb211d8197a054d91", size = 1240069, upload-time = "2026-05-15T04:50:38.221Z" },
+    { url = "https://files.pythonhosted.org/packages/aa/90/28d7f154888610aa9237e541986beb62b479df29d193a5a0617dbb1514d0/tiktoken-0.13.0-cp312-cp312-win_amd64.whl", hash = "sha256:8f2d16e7a7c783ad81f36e457d046d1f1c8af70b22aec8a13238efe531977c41", size = 874748, upload-time = "2026-05-15T04:50:39.587Z" },
+]
+
 [[package]]
 name = "tokenizers"
 version = "0.22.2"
@@ -2024,7 +2193,7 @@ wheels = [
 
 [[package]]
 name = "transformers"
-version = "5.1.0"
+version = "5.5.4"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "huggingface-hub" },
@@ -2035,11 +2204,11 @@ dependencies = [
     { name = "safetensors" },
     { name = "tokenizers" },
     { name = "tqdm" },
-    { name = "typer-slim" },
+    { name = "typer" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/c9/1d/a7d91500a6c02ec76058bc9e65fcdec1bdb8882854dec8e4adf12d0aa8b0/transformers-5.1.0.tar.gz", hash = "sha256:c60d6180e5845ea1b4eed38d7d1b06fcc4cc341c6b7fa5c1dc767d7e25fe0139", size = 8531810, upload-time = "2026-02-05T15:41:42.932Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/a5/1e/1e244ab2ab50a863e6b52cc55761910567fa532b69a6740f6e99c5fdbd98/transformers-5.5.4.tar.gz", hash = "sha256:2e67cadba81fc7608cc07c4dd54f524820bc3d95b1cabd0ef3db7733c4f8b82e", size = 8227649, upload-time = "2026-04-13T16:55:55.181Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/b7/66/57042d4b0f1ede8046d7ae6409bf3640df996e9cbc3fe20467aa29badc54/transformers-5.1.0-py3-none-any.whl", hash = "sha256:de534b50c9b2ce6217fc56421075a1734241fb40704fdc90f50f6a08fc533d59", size = 10276537, upload-time = "2026-02-05T15:41:40.358Z" },
+    { url = "https://files.pythonhosted.org/packages/29/fb/162a66789c65e5afa3b051309240c26bf37fbc8fea285b4546ae747995a2/transformers-5.5.4-py3-none-any.whl", hash = "sha256:0bd6281b82966fe5a7a16f553ea517a9db1dee6284d7cb224dfd88fc0dd1c167", size = 10236696, upload-time = "2026-04-13T16:55:51.497Z" },
 ]
 
 [[package]]
@@ -2066,18 +2235,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/4a/91/48db081e7a63bb37284f9fbcefda7c44c277b18b0e13fbc36ea2335b71e6/typer-0.24.1-py3-none-any.whl", hash = "sha256:112c1f0ce578bfb4cab9ffdabc68f031416ebcc216536611ba21f04e9aa84c9e", size = 56085, upload-time = "2026-02-21T16:54:41.616Z" },
 ]
 
-[[package]]
-name = "typer-slim"
-version = "0.24.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "typer" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/a7/a7/e6aecc4b4eb59598829a3b5076a93aff291b4fdaa2ded25efc4e1f4d219c/typer_slim-0.24.0.tar.gz", hash = "sha256:f0ed36127183f52ae6ced2ecb2521789995992c521a46083bfcdbb652d22ad34", size = 4776, upload-time = "2026-02-16T22:08:51.2Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/a7/24/5480c20380dfd18cf33d14784096dca45a24eae6102e91d49a718d3b6855/typer_slim-0.24.0-py3-none-any.whl", hash = "sha256:d5d7ee1ee2834d5020c7c616ed5e0d0f29b9a4b1dd283bdebae198ec09778d0e", size = 3394, upload-time = "2026-02-16T22:08:49.92Z" },
-]
-
 [[package]]
 name = "typing-extensions"
 version = "4.15.0"
diff --git a/tests/recipes/embed/integration/conftest.py b/tests/recipes/embed/integration/conftest.py
index 7357dd368..3f5925091 100644
--- a/tests/recipes/embed/integration/conftest.py
+++ b/tests/recipes/embed/integration/conftest.py
@@ -46,7 +46,7 @@
     },
     "stage2_finetune": {
         "key_imports": ["nemo_automodel", "omegaconf"],
-        "container": "nvcr.io/nvidia/pytorch:25.12-py3",
+        "container": "nvcr.io/nvidia/nemo-automodel:26.04",
         "entry_script": "train.py",
     },
     "stage3_eval": {
diff --git a/tests/recipes/embed/test_finetune_optimizer_backends.py b/tests/recipes/embed/test_finetune_optimizer_backends.py
new file mode 100644
index 000000000..907e65390
--- /dev/null
+++ b/tests/recipes/embed/test_finetune_optimizer_backends.py
@@ -0,0 +1,60 @@
+"""Tests for embed finetune optimizer backend selection."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import pytest
+
+from nemotron.recipes.embed.stage2_finetune import train
+
+
+def _as_dict(raw_config: dict[str, Any]) -> dict[str, Any]:
+    """Stand in for Automodel ConfigNode while preserving raw config for assertions."""
+    return raw_config
+
+
+def test_auto_optimizer_uses_fused_adam_when_available(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(train, "_can_import_fused_adam", lambda: (True, None))
+    monkeypatch.setattr(train, "_can_import_flash_adamw", lambda: (True, None))
+
+    raw_config, optimizer_backend = train._load_automodel_config(train.FinetuneConfig(), _as_dict)
+
+    assert optimizer_backend == "fused_adam"
+    assert raw_config["model"]["_target_"] == (
+        "nemo_automodel._transformers.auto_model.NeMoAutoModelBiEncoder.from_pretrained"
+    )
+    assert raw_config["dataloader"]["dataset"]["n_passages"] == 5
+    assert raw_config["dataloader"]["collate_fn"]["_target_"] == (
+        "nemo_automodel.components.datasets.llm.BiEncoderCollator"
+    )
+    assert raw_config["optimizer"]["_target_"] == (
+        "transformer_engine.pytorch.optimizers.fused_adam.FusedAdam"
+    )
+    assert raw_config["optimizer"]["adam_w_mode"] is True
+    assert raw_config["optimizer"]["master_weights"] is True
+
+
+def test_flash_adamw_backend_rewrites_optimizer_config(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(train, "_can_import_fused_adam", lambda: (False, "missing TE"))
+    monkeypatch.setattr(train, "_can_import_flash_adamw", lambda: (True, None))
+
+    cfg = train.FinetuneConfig(
+        optimizer_backend="flash_adamw",
+        flash_adamw_master_weight_bits=24,
+    )
+    raw_config, optimizer_backend = train._load_automodel_config(cfg, _as_dict)
+
+    assert optimizer_backend == "flash_adamw"
+    assert raw_config["optimizer"] == {
+        "_target_": "flashoptim.FlashAdamW",
+        "lr": 5.0e-6,
+        "weight_decay": 0.01,
+        "betas": [0.9, 0.999],
+        "eps": 1.0e-8,
+        "quantize": False,
+        "compress_state_dict": False,
+        "master_weight_bits": 24,
+        "fused": True,
+    }
+    assert raw_config["model"]["torch_dtype"] == "bfloat16"

From 01b248fbfde003057383343255f2c58aa519cb70 Mon Sep 17 00:00:00 2001
From: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>
Date: Wed, 20 May 2026 13:30:46 +0100
Subject: [PATCH 2/3] fix(embed): align biencoder config with automodel example

Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>
---
 .../embed/stage2_finetune/biencoder_base.yaml | 23 +++++++++++--------
 .../embed/stage2_finetune/pyproject.toml      |  4 ----
 .../recipes/embed/stage2_finetune/uv.lock     |  9 +++-----
 .../embed/test_finetune_optimizer_backends.py |  8 ++++---
 4 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/src/nemotron/recipes/embed/stage2_finetune/biencoder_base.yaml b/src/nemotron/recipes/embed/stage2_finetune/biencoder_base.yaml
index cd9fe67f5..2a3d9f9c1 100644
--- a/src/nemotron/recipes/embed/stage2_finetune/biencoder_base.yaml
+++ b/src/nemotron/recipes/embed/stage2_finetune/biencoder_base.yaml
@@ -14,10 +14,14 @@
 
 # Base configuration for biencoder fine-tuning using nemo-automodel
 # This file provides defaults that will be overridden by the Nemotron CLI
-# Source: https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/biencoder/llama3_2_1b_biencoder.yaml
+# Source: https://github.com/NVIDIA-NeMo/Automodel/blob/main/examples/retrieval/bi_encoder/llama3_2_1b.yaml
+
+recipe: TrainBiEncoderRecipe
 
 seed: 42
 
+temperature: 0.02
+
 step_scheduler:
   global_batch_size: 128
   local_batch_size: 4
@@ -30,7 +34,7 @@ dist_env:
   timeout_minutes: 1
 
 model:
-  _target_: nemo_automodel._transformers.auto_model.NeMoAutoModelBiEncoder.from_pretrained
+  _target_: nemo_automodel.NeMoAutoModelBiEncoder.from_pretrained
   pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
   attn_implementation: sdpa  # Use PyTorch SDPA as fallback (works with/without flash-attn)
   pooling: avg
@@ -39,22 +43,20 @@ model:
   use_sdpa_patching: true
   torch_dtype: bfloat16
 
-temperature: 0.02
-
 tokenizer:
-  _target_: nemo_automodel._transformers.auto_tokenizer.NeMoAutoTokenizer.from_pretrained
+  _target_: nemo_automodel.NeMoAutoTokenizer.from_pretrained
   pretrained_model_name_or_path: meta-llama/Llama-3.2-1B
+  add_eos_token: false
 
 dataloader:
   _target_: torchdata.stateful_dataloader.StatefulDataLoader
   dataset:
     _target_: nemo_automodel.components.datasets.llm.make_retrieval_dataset
+    model_type: bi_encoder
     data_dir_list:
       - training_datasets/nqsh_shuffled_50k.json
-    model_type: bi_encoder
     data_type: train
     n_passages: 5
-    eval_negative_size: 4
     seed: 42
     do_shuffle: true
   collate_fn:
@@ -85,6 +87,9 @@ checkpoint:
   model_save_format: safetensors
   save_consolidated: true
 
-distributed_config:
-  _target_: nemo_automodel.components.distributed.config.FSDP2Config
+distributed:
+  strategy: fsdp2
+  dp_size: none
+  tp_size: 1
+  cp_size: 1
   sequence_parallel: false
diff --git a/src/nemotron/recipes/embed/stage2_finetune/pyproject.toml b/src/nemotron/recipes/embed/stage2_finetune/pyproject.toml
index 731324e4b..3e2bb03fc 100644
--- a/src/nemotron/recipes/embed/stage2_finetune/pyproject.toml
+++ b/src/nemotron/recipes/embed/stage2_finetune/pyproject.toml
@@ -14,10 +14,6 @@ dependencies = [
 # Note: exclude-dependencies are injected dynamically by run_uv.py for Docker/Slurm
 # This base config works for local mode where all packages need to be installed
 
-[tool.uv]
-# Override automodel's transformers pin to use 5.5.x
-override-dependencies = ["transformers>=5.5.0,<5.6.0"]
-
 # Pin Linux torch installs to cu129. CUDA 13 drivers can run cu12.x wheels via
 # NVIDIA driver backward compatibility, while CUDA 12.9 hosts cannot run cu130.
 [tool.uv.sources]
diff --git a/src/nemotron/recipes/embed/stage2_finetune/uv.lock b/src/nemotron/recipes/embed/stage2_finetune/uv.lock
index 8f5ee0168..afe69d217 100644
--- a/src/nemotron/recipes/embed/stage2_finetune/uv.lock
+++ b/src/nemotron/recipes/embed/stage2_finetune/uv.lock
@@ -8,9 +8,6 @@ resolution-markers = [
     "sys_platform == 'darwin'",
 ]
 
-[manifest]
-overrides = [{ name = "transformers", specifier = ">=5.5.0,<5.6.0" }]
-
 [[package]]
 name = "aiohappyeyeballs"
 version = "2.6.1"
@@ -2193,7 +2190,7 @@ wheels = [
 
 [[package]]
 name = "transformers"
-version = "5.5.4"
+version = "5.5.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
     { name = "huggingface-hub" },
@@ -2206,9 +2203,9 @@ dependencies = [
     { name = "tqdm" },
     { name = "typer" },
 ]
-sdist = { url = "https://files.pythonhosted.org/packages/a5/1e/1e244ab2ab50a863e6b52cc55761910567fa532b69a6740f6e99c5fdbd98/transformers-5.5.4.tar.gz", hash = "sha256:2e67cadba81fc7608cc07c4dd54f524820bc3d95b1cabd0ef3db7733c4f8b82e", size = 8227649, upload-time = "2026-04-13T16:55:55.181Z" }
+sdist = { url = "https://files.pythonhosted.org/packages/ff/9d/fb46e729b461985f41a5740167688b924a4019141e5c164bea77548d3d9e/transformers-5.5.0.tar.gz", hash = "sha256:c8db656cf51c600cd8c75f06b20ef85c72e8b8ff9abc880c5d3e8bc70e0ddcbd", size = 8237745, upload-time = "2026-04-02T16:13:08.113Z" }
 wheels = [
-    { url = "https://files.pythonhosted.org/packages/29/fb/162a66789c65e5afa3b051309240c26bf37fbc8fea285b4546ae747995a2/transformers-5.5.4-py3-none-any.whl", hash = "sha256:0bd6281b82966fe5a7a16f553ea517a9db1dee6284d7cb224dfd88fc0dd1c167", size = 10236696, upload-time = "2026-04-13T16:55:51.497Z" },
+    { url = "https://files.pythonhosted.org/packages/e7/28/35f7411ff80a3640c1f4fc907dcbb6a65061ebb82f66950e38bfc9f7f740/transformers-5.5.0-py3-none-any.whl", hash = "sha256:821a9ff0961abbb29eb1eb686d78df1c85929fdf213a3fe49dc6bd94f9efa944", size = 10245591, upload-time = "2026-04-02T16:13:03.462Z" },
 ]
 
 [[package]]
diff --git a/tests/recipes/embed/test_finetune_optimizer_backends.py b/tests/recipes/embed/test_finetune_optimizer_backends.py
index 907e65390..a0a098605 100644
--- a/tests/recipes/embed/test_finetune_optimizer_backends.py
+++ b/tests/recipes/embed/test_finetune_optimizer_backends.py
@@ -21,13 +21,15 @@ def test_auto_optimizer_uses_fused_adam_when_available(monkeypatch: pytest.Monke
     raw_config, optimizer_backend = train._load_automodel_config(train.FinetuneConfig(), _as_dict)
 
     assert optimizer_backend == "fused_adam"
-    assert raw_config["model"]["_target_"] == (
-        "nemo_automodel._transformers.auto_model.NeMoAutoModelBiEncoder.from_pretrained"
-    )
+    assert raw_config["recipe"] == "TrainBiEncoderRecipe"
+    assert raw_config["model"]["_target_"] == "nemo_automodel.NeMoAutoModelBiEncoder.from_pretrained"
+    assert raw_config["tokenizer"]["_target_"] == "nemo_automodel.NeMoAutoTokenizer.from_pretrained"
+    assert raw_config["tokenizer"]["add_eos_token"] is False
     assert raw_config["dataloader"]["dataset"]["n_passages"] == 5
     assert raw_config["dataloader"]["collate_fn"]["_target_"] == (
         "nemo_automodel.components.datasets.llm.BiEncoderCollator"
     )
+    assert raw_config["distributed"]["strategy"] == "fsdp2"
     assert raw_config["optimizer"]["_target_"] == (
         "transformer_engine.pytorch.optimizers.fused_adam.FusedAdam"
     )

From 011ceb7b85a1ec544d171f2b178abd6ace6ca908 Mon Sep 17 00:00:00 2001
From: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>
Date: Wed, 20 May 2026 14:42:01 +0100
Subject: [PATCH 3/3] fix(embed): trust remote code for default model

Signed-off-by: Oliver Holworthy <1216955+oliverholworthy@users.noreply.github.com>
---
 src/nemotron/recipes/embed/README.md                        | 1 +
 .../recipes/embed/stage2_finetune/config/default.yaml       | 1 +
 src/nemotron/recipes/embed/stage2_finetune/train.py         | 6 ++++++
 3 files changed, 8 insertions(+)

diff --git a/src/nemotron/recipes/embed/README.md b/src/nemotron/recipes/embed/README.md
index b8f7357d5..3060477b4 100644
--- a/src/nemotron/recipes/embed/README.md
+++ b/src/nemotron/recipes/embed/README.md
@@ -448,6 +448,7 @@ test_ratio: 0.1                # Test split (10%)
 **Stage 2: Finetune**
 ```yaml
 base_model: nvidia/llama-nemotron-embed-1b-v2
+trust_remote_code: true
 num_epochs: 3
 global_batch_size: 128
 learning_rate: 1.0e-5
diff --git a/src/nemotron/recipes/embed/stage2_finetune/config/default.yaml b/src/nemotron/recipes/embed/stage2_finetune/config/default.yaml
index 32ff77660..2e374925b 100644
--- a/src/nemotron/recipes/embed/stage2_finetune/config/default.yaml
+++ b/src/nemotron/recipes/embed/stage2_finetune/config/default.yaml
@@ -18,6 +18,7 @@ run:
 
 # Base embedding model to fine-tune
 base_model: nvidia/llama-nemotron-embed-1b-v2
+trust_remote_code: true
 
 # Path to training data (output from stage1_data_prep)
 train_data_path: ./output/embed/stage1_data_prep/train_mined.automodel_unrolled.json
diff --git a/src/nemotron/recipes/embed/stage2_finetune/train.py b/src/nemotron/recipes/embed/stage2_finetune/train.py
index edc03a7f0..69f026b99 100644
--- a/src/nemotron/recipes/embed/stage2_finetune/train.py
+++ b/src/nemotron/recipes/embed/stage2_finetune/train.py
@@ -83,6 +83,10 @@ class FinetuneConfig(RecipeSettings):
         default="nvidia/llama-nemotron-embed-1b-v2",
         description="Base embedding model to fine-tune.",
     )
+    trust_remote_code: bool = Field(
+        default=True,
+        description="Allow Hugging Face custom model code. Required by the default Nemotron Embed model.",
+    )
 
     # Data paths
     train_data_path: Path = Field(
@@ -379,6 +383,8 @@ def run_finetune(cfg: FinetuneConfig) -> Path:
     # Model settings
     automodel_cfg.model.pretrained_model_name_or_path = cfg.base_model
     automodel_cfg.tokenizer.pretrained_model_name_or_path = cfg.base_model
+    automodel_cfg.model.trust_remote_code = cfg.trust_remote_code
+    automodel_cfg.tokenizer.trust_remote_code = cfg.trust_remote_code
     # Auto-detect attention implementation if not explicitly set
     if cfg.attn_implementation is not None:
         attn_impl = cfg.attn_implementation