diff --git a/examples/diffusion/finetune/flux_t2i_flow.yaml b/examples/diffusion/finetune/flux_t2i_flow.yaml
index 4d64550a5e..cce14da582 100644
--- a/examples/diffusion/finetune/flux_t2i_flow.yaml
+++ b/examples/diffusion/finetune/flux_t2i_flow.yaml
@@ -3,6 +3,10 @@ model:
   mode: "finetune"
   cache_dir: null
   attention_backend: "flash"
+  transformer_engine_fp8: true
+  transformer_engine_fp8_recipe: "delayed"
+  transformer_engine_fp8_amax_history_len: 1024
+  transformer_engine_fp8_amax_compute_algo: "max"
 
 optim:
   learning_rate: 2e-5
@@ -10,7 +14,8 @@ optim:
   optimizer:
     weight_decay: 0.01
     betas: [0.9, 0.999]
-    foreach: true
+    foreach: false
+    fused: true
 
 performance:
   check_loss: false
@@ -25,9 +30,10 @@ fsdp:
   activation_checkpointing: false
   cpu_offload: false
   defer_fsdp_grad_sync: true
+  enable_compile: false
   enable_fsdp2_prefetch: true
-  fsdp2_backward_prefetch_depth: 2
-  fsdp2_forward_prefetch_depth: 1
+  fsdp2_backward_prefetch_depth: 3
+  fsdp2_forward_prefetch_depth: 2
 
 flow_matching:
   adapter_type: "flux"
@@ -49,8 +55,8 @@ flow_matching:
 
 step_scheduler:
   num_epochs: 5000
-  local_batch_size: 1
-  global_batch_size: 8
+  local_batch_size: 4
+  global_batch_size: 32
   ckpt_every_steps: 500
   log_every: 1
   # max_steps: null  # Set to limit training to a specific number of steps
diff --git a/examples/diffusion/finetune/flux_t2i_flow_lora.yaml b/examples/diffusion/finetune/flux_t2i_flow_lora.yaml
index 15cc56f0be..7f417a9205 100644
--- a/examples/diffusion/finetune/flux_t2i_flow_lora.yaml
+++ b/examples/diffusion/finetune/flux_t2i_flow_lora.yaml
@@ -17,6 +17,7 @@ model:
   mode: "finetune"
   cache_dir: null
   attention_backend: "flash"
+  transformer_engine_fp8: false
 
 # ── LoRA / PEFT configuration ─────────────────────────────────────────────────
 # Remove this section for full fine-tune.
@@ -45,22 +46,18 @@ optim:
   optimizer:
     weight_decay: 0.01
     betas: [0.9, 0.999]
+    foreach: false
+    fused: true
 
 lr_scheduler:
   lr_decay_style: cosine
   lr_warmup_steps: 100
   min_lr: 1e-6
 
-# FSDP2 with param_dtype=None when LoRA enabled:
-#   base weights (bf16): sharded, no grad comms
-#   LoRA weights (fp32): sharded, grad allreduced by FSDP2
-fsdp:
-  dp_size: 8
-  tp_size: 1
-  cp_size: 1
-  pp_size: 1
+# Best-performing Flux LoRA config from the local perf sweep uses replicated DDP.
+ddp:
+  backend: "nccl"
   activation_checkpointing: false
-  cpu_offload: false
 
 flow_matching:
   adapter_type: "flux"
@@ -82,8 +79,8 @@ flow_matching:
 
 step_scheduler:
   num_epochs: 5000
-  local_batch_size: 1
-  global_batch_size: 8
+  local_batch_size: 6
+  global_batch_size: 48
   ckpt_every_steps: 500
   log_every: 1
   # max_steps: null  # Set to limit training to a specific number of steps
@@ -111,4 +108,4 @@ checkpoint:
 
 ci:
   recipe_owner: pthombre
-  time: "00:30:00"
\ No newline at end of file
+  time: "00:30:00"
diff --git a/examples/diffusion/generate/generate.py b/examples/diffusion/generate/generate.py
index 6a164041bb..45e2e9004f 100644
--- a/examples/diffusion/generate/generate.py
+++ b/examples/diffusion/generate/generate.py
@@ -266,10 +266,30 @@ def load_lora_weights_into_pipeline(pipe, cfg):
         raise FileNotFoundError(f"LoRA weights directory not found: {lora_path}")
 
     with open(lora_path / "adapter_config.json") as f:
-        peft_config = PeftConfig.from_dict(json.load(f))
-    apply_lora_to_linear_modules(pipe.transformer, peft_config, skip_freeze=True)
+        peft_config_dict = json.load(f)
+    if "dim" not in peft_config_dict and "r" in peft_config_dict:
+        peft_config_dict["dim"] = peft_config_dict["r"]
+    if "alpha" not in peft_config_dict and "lora_alpha" in peft_config_dict:
+        peft_config_dict["alpha"] = peft_config_dict["lora_alpha"]
+    if "dropout" not in peft_config_dict and "lora_dropout" in peft_config_dict:
+        peft_config_dict["dropout"] = peft_config_dict["lora_dropout"]
+
+    peft_config = PeftConfig.from_dict(peft_config_dict)
+    num_lora_modules = apply_lora_to_linear_modules(pipe.transformer, peft_config, skip_freeze=True)
+    if num_lora_modules == 0:
+        raise RuntimeError(f"No transformer modules matched LoRA config from {lora_path / 'adapter_config.json'}")
+
     state_dict = load_file(lora_path / "adapter_model.safetensors", device="cuda")
-    pipe.transformer.load_state_dict(state_dict, strict=False)
+    state_dict = {key.removeprefix("base_model.model."): value for key, value in state_dict.items()}
+    load_result = pipe.transformer.load_state_dict(state_dict, strict=False)
+    missing_lora_keys = sorted(key for key in load_result.missing_keys if ".lora_" in key)
+    unexpected_lora_keys = sorted(key for key in load_result.unexpected_keys if ".lora_" in key)
+    if missing_lora_keys or unexpected_lora_keys:
+        raise RuntimeError(
+            "LoRA checkpoint did not load cleanly. "
+            f"missing_lora_keys={missing_lora_keys[:10]}, unexpected_lora_keys={unexpected_lora_keys[:10]}"
+        )
+    logger.info("Loaded LoRA adapter from %s (%d modules, %d tensors)", lora_path, num_lora_modules, len(state_dict))
 
 
 def _load_sharded_fsdp_checkpoint(transformer, sharded_dir, torch_dtype=torch.bfloat16):
diff --git a/nemo_automodel/_diffusers/auto_diffusion_pipeline.py b/nemo_automodel/_diffusers/auto_diffusion_pipeline.py
index 3ad3a0d495..256be540ce 100644
--- a/nemo_automodel/_diffusers/auto_diffusion_pipeline.py
+++ b/nemo_automodel/_diffusers/auto_diffusion_pipeline.py
@@ -57,6 +57,7 @@
     HunyuanParallelizationStrategy,
     WanParallelizationStrategy,
 )
+from nemo_automodel.shared.import_utils import safe_import_te
 from nemo_automodel.shared.utils import dtype_from_str
 
 # diffusers is an optional dependency
@@ -204,6 +205,132 @@ def _ensure_params_trainable(module: nn.Module, module_name: Optional[str] = Non
     return num_trainable_parameters
 
 
+def _is_fp8_training_safe_linear(name: str, linear: nn.Linear) -> bool:
+    """Return whether a Linear has shapes that are safe for FP8 training kernels."""
+    if linear.weight.shape[0] % 16 != 0 or linear.weight.shape[1] % 16 != 0:
+        return False
+    if name.startswith(("time_text_embed.", "norm_out.")):
+        return False
+    if name.endswith(".linear") and any(norm_name in name for norm_name in (".norm.", ".norm1.", ".norm1_context.")):
+        return False
+    return True
+
+
+def _replace_linear_with_transformer_engine(
+    module: nn.Module,
+    module_name: str,
+    *,
+    fp8_safe_only: bool = False,
+) -> int:
+    """Replace torch Linear modules with Transformer Engine Linear modules."""
+    has_te, transformer_engine = safe_import_te()
+    if not has_te:
+        raise ImportError("model.transformer_engine_linear=true requires transformer_engine.pytorch to be importable")
+
+    te_linear_cls = transformer_engine.pytorch.Linear
+    converted = 0
+    skipped = 0
+
+    def replace_children(parent: nn.Module, prefix: str = "") -> None:
+        nonlocal converted, skipped
+        for child_name, child in list(parent.named_children()):
+            child_fqn = f"{prefix}.{child_name}" if prefix else child_name
+            if isinstance(child, nn.Linear):
+                if fp8_safe_only and not _is_fp8_training_safe_linear(child_fqn, child):
+                    skipped += 1
+                    logger.info(
+                        "[TE] Keeping %s.%s as torch.nn.Linear for FP8 shape compatibility; weight shape=%s",
+                        module_name,
+                        child_fqn,
+                        tuple(child.weight.shape),
+                    )
+                    continue
+                te_linear = te_linear_cls(
+                    child.in_features,
+                    child.out_features,
+                    bias=child.bias is not None,
+                    device=child.weight.device,
+                    params_dtype=child.weight.dtype,
+                )
+                te_linear.train(child.training)
+                with torch.no_grad():
+                    te_linear.weight.copy_(child.weight)
+                    te_linear.weight.requires_grad_(child.weight.requires_grad)
+                    if child.bias is not None:
+                        te_linear.bias.copy_(child.bias)
+                        te_linear.bias.requires_grad_(child.bias.requires_grad)
+                setattr(parent, child_name, te_linear)
+                converted += 1
+            elif not isinstance(child, te_linear_cls):
+                replace_children(child, child_fqn)
+
+    replace_children(module)
+    logger.info(
+        "[TE] Replaced %d torch.nn.Linear modules with Transformer Engine Linear in %s; skipped=%d fp8_safe_only=%s",
+        converted,
+        module_name,
+        skipped,
+        fp8_safe_only,
+    )
+    return converted
+
+
+def _remove_child_module(parent: nn.Module, child_name: str) -> int:
+    """Remove a child module and return the number of parameters it owned."""
+    child = getattr(parent, child_name, None)
+    if not isinstance(child, nn.Module):
+        return 0
+
+    parameter_count = sum(parameter.numel() for parameter in child.parameters())
+    delattr(parent, child_name)
+    return parameter_count
+
+
+def _compact_fused_qkv_projection_modules(module: nn.Module, module_name: str) -> tuple[int, int]:
+    """Remove original projection modules that are unused after Diffusers QKV fusion."""
+    removed_modules = 0
+    removed_parameters = 0
+
+    for child in list(module.modules()):
+        if not getattr(child, "fused_projections", False):
+            continue
+
+        if hasattr(child, "to_qkv"):
+            for projection_name in ("to_q", "to_k", "to_v"):
+                parameter_count = _remove_child_module(child, projection_name)
+                if parameter_count:
+                    removed_modules += 1
+                    removed_parameters += parameter_count
+
+        if hasattr(child, "to_added_qkv"):
+            for projection_name in ("add_q_proj", "add_k_proj", "add_v_proj"):
+                parameter_count = _remove_child_module(child, projection_name)
+                if parameter_count:
+                    removed_modules += 1
+                    removed_parameters += parameter_count
+
+    logger.info(
+        "[QKV] Removed %d original projection modules with %s parameters after fused QKV setup in %s",
+        removed_modules,
+        f"{removed_parameters:,}",
+        module_name,
+    )
+    return removed_modules, removed_parameters
+
+
+def _fuse_transformer_qkv_projections(module: nn.Module, module_name: str, *, compact: bool = False) -> int:
+    """Fuse Diffusers attention QKV projections when the transformer exposes that hook."""
+    if not hasattr(module, "fuse_qkv_projections"):
+        raise AttributeError(f"model.fuse_qkv_projections=true requires {module_name} to expose fuse_qkv_projections()")
+
+    module.fuse_qkv_projections()
+    fused = sum(1 for child in module.modules() if getattr(child, "fused_projections", False))
+    logger.info("[QKV] Fused QKV projections for %d attention modules in %s", fused, module_name)
+    if compact:
+        _compact_fused_qkv_projection_modules(module, module_name)
+    return fused
+
+
 def _create_parallel_manager(manager_args: Dict[str, Any]) -> ParallelManager:
     """
     Factory function to create the appropriate parallel manager based on config.
@@ -309,6 +436,8 @@ def _apply_parallelization(
         manager = _create_parallel_manager(manager_args)
         created_managers[comp_name] = manager
         parallel_module = manager.parallelize(comp_module)
+        if hasattr(manager, "maybe_compile"):
+            manager.maybe_compile(parallel_module)
         setattr(pipe, comp_name, parallel_module)
 
     return created_managers
@@ -373,6 +502,10 @@ def from_pretrained(
         components_to_load: Optional[Iterable[str]] = None,
         peft_cfg=None,
         model_type=None,
+        transformer_engine_linear: bool = False,
+        transformer_engine_fp8_safe_only: bool = False,
+        fuse_qkv_projections: bool = False,
+        compact_fused_qkv_projections: bool = False,
         **kwargs,
     ) -> Tuple[DiffusionPipeline, Dict[str, ParallelManager]]:
         """
@@ -396,6 +529,10 @@ def from_pretrained(
                 before _apply_parallelization() (FSDP2 wrapping). Base weights
                 are frozen after FSDP2; LoRA params are collected pre-FSDP2 and stored on pipe.
             model_type: "flux" | "wan" | "hunyuan". Required when peft_cfg is provided.
+            transformer_engine_linear: Whether to replace torch.nn.Linear modules in the transformer with TE Linear.
+            transformer_engine_fp8_safe_only: Whether to skip TE Linear conversion for known FP8-incompatible modules.
+            fuse_qkv_projections: Whether to call Diffusers QKV projection fusion on the transformer.
+            compact_fused_qkv_projections: Whether to remove original projection modules after QKV fusion.
             **kwargs: Additional arguments passed to DiffusionPipeline.from_pretrained
 
         Returns:
@@ -429,6 +566,23 @@ def from_pretrained(
                     logger.info("[INFO] Moving module: %s to device/dtype", name)
                     _move_module_to_device(module, dev, torch_dtype)
 
+        if compact_fused_qkv_projections and not fuse_qkv_projections:
+            raise ValueError("model.compact_fused_qkv_projections=true requires model.fuse_qkv_projections=true")
+
+        if fuse_qkv_projections:
+            for name, module in _iter_pipeline_modules(pipe):
+                if name == "transformer" and (not components_to_load or name in components_to_load):
+                    _fuse_transformer_qkv_projections(module, name, compact=compact_fused_qkv_projections)
+
+        if transformer_engine_linear:
+            for name, module in _iter_pipeline_modules(pipe):
+                if name == "transformer" and (not components_to_load or name in components_to_load):
+                    _replace_linear_with_transformer_engine(
+                        module,
+                        name,
+                        fp8_safe_only=transformer_engine_fp8_safe_only,
+                    )
+
         if peft_cfg is not None:
             # ── LoRA path ─────────────────────────────────────────────────────
             # LoRA injection MUST run before _apply_parallelization (FSDP2).
@@ -483,20 +637,31 @@ def from_pretrained(
                         logger.info("[INFO] Ensuring params trainable: %s", name)
                         _ensure_params_trainable(module, module_name=name)
 
+        if peft_cfg is not None:
+            transformer_parallel_args = (parallel_scheme or {}).get("transformer", {})
+            manager_type = str(transformer_parallel_args.get("_manager_type", "fsdp2")).lower()
+            if manager_type == "ddp":
+                for param_name, param in pipe.transformer.named_parameters():
+                    if "lora_" not in param_name and param.requires_grad:
+                        param.requires_grad_(False)
+                logger.info("[LoRA] Froze base weights before DDP wrapping")
+
         # Apply parallelization (FSDP2 or DDP)
-        # LoRA: all params are trainable when fully_shard() runs so FSDP2 sets
-        # up gradient reduction for lora_A/lora_B correctly. Freeze happens below.
+        # FSDP2 LoRA: all params are trainable when fully_shard() runs so FSDP2
+        # sets up gradient reduction for lora_A/lora_B correctly. Freeze happens below.
+        # DDP LoRA: base weights are frozen before wrapping so DDP only reduces LoRA gradients.
         created_managers = _apply_parallelization(pipe, parallel_scheme)
 
-        # Freeze base weights AFTER FSDP2 wrapping — mirrors the LLM pattern in
+        # Freeze base weights after FSDP2 wrapping — mirrors the LLM pattern in
         # nemo_automodel/_transformers/infrastructure.py lines 513-518.
         # FSDP2 must see all params as trainable during fully_shard(); freezing
         # before wrapping causes gradient reduction to malfunction for LoRA params.
+        # For DDP this is harmless because the base was already frozen before wrapping.
         if peft_cfg is not None:
             for name, param in pipe.transformer.named_parameters():
                 if "lora_" not in name and param.requires_grad:
                     param.requires_grad_(False)
-            logger.info("[LoRA] Froze base weights after FSDP2 wrapping")
+            logger.info("[LoRA] Froze base weights after parallelization")
 
         return pipe, created_managers
 
@@ -510,6 +675,10 @@ def from_config(
         parallel_scheme: Optional[Dict[str, Dict[str, Any]]] = None,
         move_to_device: bool = True,
         components_to_load: Optional[Iterable[str]] = None,
+        transformer_engine_linear: bool = False,
+        transformer_engine_fp8_safe_only: bool = False,
+        fuse_qkv_projections: bool = False,
+        compact_fused_qkv_projections: bool = False,
         **kwargs,
     ) -> Tuple["NeMoAutoDiffusionPipeline", Dict[str, ParallelManager]]:
         """
@@ -531,6 +700,10 @@ def from_config(
             parallel_scheme: Dict mapping component names to parallel manager kwargs
             move_to_device: Whether to move modules to device
             components_to_load: Which components to process (default: all)
+            transformer_engine_linear: Whether to replace torch.nn.Linear modules in the transformer with TE Linear.
+            transformer_engine_fp8_safe_only: Whether to skip TE Linear conversion for known FP8-incompatible modules.
+            fuse_qkv_projections: Whether to call Diffusers QKV projection fusion on the transformer.
+            compact_fused_qkv_projections: Whether to remove original projection modules after QKV fusion.
             **kwargs: Additional arguments
 
         Returns:
@@ -588,9 +761,23 @@ def from_config(
                 transformer = transformer.to(dev)
             pipe = cls(transformer=transformer)
 
+        if compact_fused_qkv_projections and not fuse_qkv_projections:
+            raise ValueError("model.compact_fused_qkv_projections=true requires model.fuse_qkv_projections=true")
+
+        if fuse_qkv_projections:
+            for name, module in _iter_pipeline_modules(pipe):
+                if name == "transformer" and (not components_to_load or name in components_to_load):
+                    _fuse_transformer_qkv_projections(module, name, compact=compact_fused_qkv_projections)
+
         # Make parameters trainable (always true for from_config / pretraining)
         for name, module in _iter_pipeline_modules(pipe):
             if not components_to_load or name in components_to_load:
+                if transformer_engine_linear and name == "transformer":
+                    _replace_linear_with_transformer_engine(
+                        module,
+                        name,
+                        fp8_safe_only=transformer_engine_fp8_safe_only,
+                    )
                 _ensure_params_trainable(module, module_name=name)
 
         # Apply parallelization (FSDP2 or DDP)
diff --git a/nemo_automodel/components/distributed/parallelizer.py b/nemo_automodel/components/distributed/parallelizer.py
index 93db0b6d54..dea1da3185 100644
--- a/nemo_automodel/components/distributed/parallelizer.py
+++ b/nemo_automodel/components/distributed/parallelizer.py
@@ -743,24 +743,47 @@ def _apply_per_layer_compile(model: nn.Module) -> None:
 
     _patch_dtensor_spec_hash_for_symint()
 
+    compiled_count = 0
+    compiled_modules: set[int] = set()
+
+    def _compile_module_list(module_list: nn.ModuleList | nn.ModuleDict) -> None:
+        nonlocal compiled_count
+        # PP converts model.model.layers from nn.ModuleList to nn.ModuleDict (str keys).
+        # enumerate(nn.ModuleDict) yields string keys, not modules -- use .items() instead.
+        items = module_list.items() if isinstance(module_list, nn.ModuleDict) else enumerate(module_list)
+        for _, layer in items:
+            # Unwrap any layer-level checkpoint wrapper (PP path) to reach the actual decoder layer.
+            actual_layer = layer._checkpoint_wrapped_module if isinstance(layer, CheckpointWrapper) else layer
+            module_id = id(actual_layer)
+            if module_id in compiled_modules:
+                continue
+            actual_layer.compile()
+            compiled_modules.add(module_id)
+            compiled_count += 1
+
+    module_lists = []
     if hasattr(model, "model") and hasattr(model.model, "layers"):
-        module_list = model.model.layers
-    elif hasattr(model, "layers"):
-        module_list = model.layers
+        module_lists.append(model.model.layers)
+    if hasattr(model, "layers"):
+        module_lists.append(model.layers)
+    for attr_name in ("transformer_blocks", "single_transformer_blocks"):
+        module_list = getattr(model, attr_name, None)
+        if isinstance(module_list, (nn.ModuleList, nn.ModuleDict)):
+            module_lists.append(module_list)
+
+    if module_lists:
+        for module_list in module_lists:
+            _compile_module_list(module_list)
     else:
-        logger.warning("_apply_per_layer_compile: cannot find transformer layers, skipping")
-        return
-
-    # PP converts model.model.layers from nn.ModuleList to nn.ModuleDict (str keys).
-    # enumerate(nn.ModuleDict) yields string keys, not modules -- use .items() instead.
-    items = module_list.items() if isinstance(module_list, nn.ModuleDict) else enumerate(module_list)
-
-    compiled_count = 0
-    for _, layer in items:
-        # Unwrap any layer-level checkpoint wrapper (PP path) to reach the actual decoder layer.
-        actual_layer = layer._checkpoint_wrapped_module if isinstance(layer, CheckpointWrapper) else layer
-        actual_layer.compile()
-        compiled_count += 1
+        logger.warning("_apply_per_layer_compile: using heuristic layer extraction")
+        for layer in _extract_model_layers(model):
+            actual_layer = layer._checkpoint_wrapped_module if isinstance(layer, CheckpointWrapper) else layer
+            module_id = id(actual_layer)
+            if module_id in compiled_modules:
+                continue
+            actual_layer.compile()
+            compiled_modules.add(module_id)
+            compiled_count += 1
 
     logger.info("Per-layer torch.compile applied to %d decoder layers", compiled_count)
 
diff --git a/nemo_automodel/recipes/diffusion/train.py b/nemo_automodel/recipes/diffusion/train.py
index 99ebcca782..38e38c128b 100644
--- a/nemo_automodel/recipes/diffusion/train.py
+++ b/nemo_automodel/recipes/diffusion/train.py
@@ -14,9 +14,11 @@
 
 from __future__ import annotations
 
+import inspect
 import logging
 import os
 import time
+from contextlib import nullcontext
 from math import ceil
 from typing import Any, Dict, Optional
 
@@ -44,8 +46,19 @@
 )
 from nemo_automodel.recipes.base_recipe import BaseRecipe
 from nemo_automodel.recipes.llm.train_ft import build_distributed, build_wandb
+from nemo_automodel.shared.import_utils import safe_import_from
 from nemo_automodel.shared.utils import dtype_from_str
 
+_OPTIMIZER_DEFAULT_TARGET = "torch.optim.AdamW"
+_TORCH_DTYPE_ALIASES = {
+    "float32": torch.float32,
+    "fp32": torch.float32,
+    "bfloat16": torch.bfloat16,
+    "bf16": torch.bfloat16,
+    "float16": torch.float16,
+    "fp16": torch.float16,
+}
+
 
 def _resolve_optimizer_dtype_strings(optimizer_cfg: Any) -> None:
     """Resolve dtype strings in optimizer config objects in place."""
@@ -96,15 +109,35 @@ def _build_optimizer(
     learning_rate: float,
     is_peft: bool = False,
 ) -> torch.optim.Optimizer:
-    """Build optimizer from config, falling back to AdamW for legacy dict configs."""
+    """Build optimizer from config, falling back to AdamW for legacy configs."""
     optimizer_cfg = optimizer_cfg or {}
     if isinstance(optimizer_cfg, dict) and "_target_" in optimizer_cfg:
         optimizer_cfg = ConfigNode(optimizer_cfg)
 
-    if hasattr(optimizer_cfg, "instantiate") and hasattr(optimizer_cfg, "_target_"):
+    if hasattr(optimizer_cfg, "_target_"):
         _resolve_optimizer_dtype_strings(optimizer_cfg)
-        optimizer = optimizer_cfg.instantiate(params=trainable_params, lr=learning_rate)
-        logging.info("[INFO] Optimizer: %s, lr=%s", optimizer.__class__.__name__, learning_rate)
+        optimizer_target = optimizer_cfg._target_
+        optimizer_target_name = _get_optimizer_target_name(optimizer_target)
+        optimizer_cls = _resolve_optimizer_class(optimizer_target)
+        optimizer_dict = optimizer_cfg.to_dict() if hasattr(optimizer_cfg, "to_dict") else dict(optimizer_cfg)
+
+        optimizer_kwargs = {"lr": learning_rate}
+        for key, value in optimizer_dict.items():
+            if key in {"_target_", "lr", "learning_rate"}:
+                continue
+            if value is not None:
+                optimizer_kwargs[key] = _normalize_optimizer_value(value)
+
+        if (
+            optimizer_cls is torch.optim.AdamW
+            and optimizer_kwargs.get("foreach", False)
+            and optimizer_kwargs.get("fused", False)
+        ):
+            raise ValueError("torch.optim.AdamW does not support foreach=True and fused=True at the same time")
+        optimizer_kwargs = _filter_optimizer_kwargs(optimizer_target_name, optimizer_cls, optimizer_kwargs)
+        optimizer = optimizer_cls(trainable_params, **optimizer_kwargs)
+        logging.info("[INFO] Optimizer target: %s", optimizer_target_name)
+        logging.info("[INFO] Optimizer config: %s", optimizer_kwargs)
         warn_if_torch_adam_with_bf16_params(
             optimizer=optimizer,
             optimizer_cfg=optimizer_cfg,
@@ -186,6 +219,133 @@ def _calculate_throughput_metrics(
     }
 
 
+def _normalize_optimizer_value(value: Any) -> Any:
+    """Convert CLI-friendly optimizer scalar values into Python objects."""
+    if isinstance(value, str):
+        normalized = value.removeprefix("torch.").lower()
+        return _TORCH_DTYPE_ALIASES.get(normalized, value)
+    return value
+
+
+def _get_optimizer_target_name(target: Any) -> str:
+    """Return a stable display name for an optimizer target."""
+    if isinstance(target, str):
+        return target
+    module_name = getattr(target, "__module__", None)
+    qualname = getattr(target, "__qualname__", None)
+    if module_name and qualname:
+        return f"{module_name}.{qualname}"
+    return repr(target)
+
+
+def _resolve_optimizer_class(target: Any) -> Any:
+    """Resolve an optimizer class from a fully qualified `_target_` string."""
+    if not isinstance(target, str):
+        if callable(target):
+            return target
+        raise ValueError(f"Optimizer target must be a fully qualified import path or callable, got {target!r}.")
+
+    if target == _OPTIMIZER_DEFAULT_TARGET:
+        return torch.optim.AdamW
+
+    module_name, _, symbol_name = target.rpartition(".")
+    if not module_name or not symbol_name:
+        raise ValueError(
+            f"Optimizer target must be a fully qualified import path, got {target!r}. "
+            f"Example: {_OPTIMIZER_DEFAULT_TARGET!r}."
+        )
+
+    available, optimizer_cls = safe_import_from(
+        module_name,
+        symbol_name,
+        msg=f"Optimizer target {target!r} could not be imported",
+    )
+    if not available:
+        raise ImportError(f"Optimizer target {target!r} could not be imported")
+    return optimizer_cls
+
+
+def _filter_optimizer_kwargs(target: str, optimizer_cls: Any, optimizer_kwargs: Dict[str, Any]) -> Dict[str, Any]:
+    """Drop kwargs unsupported by an opt-in optimizer target."""
+    try:
+        signature = inspect.signature(optimizer_cls)
+    except (TypeError, ValueError):
+        logging.info("[INFO] Could not inspect optimizer target %s; passing all optimizer kwargs", target)
+        return optimizer_kwargs
+
+    parameters = signature.parameters
+    if any(parameter.kind == inspect.Parameter.VAR_KEYWORD for parameter in parameters.values()):
+        return optimizer_kwargs
+
+    accepted = {
+        name
+        for name, parameter in parameters.items()
+        if parameter.kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.KEYWORD_ONLY)
+    }
+    accepted.discard("params")
+
+    filtered_kwargs = {key: value for key, value in optimizer_kwargs.items() if key in accepted}
+    ignored_kwargs = sorted(set(optimizer_kwargs) - set(filtered_kwargs))
+    if ignored_kwargs:
+        logging.info("[INFO] Optimizer target %s does not accept kwargs %s; ignoring them", target, ignored_kwargs)
+    return filtered_kwargs
+
+
+def _build_transformer_engine_fp8_recipe(
+    recipe_name: str,
+    *,
+    amax_history_len: int,
+    amax_compute_algo: str,
+) -> Any:
+    """Build a Transformer Engine FP8 recipe from CLI-friendly config values."""
+    normalized_recipe_name = recipe_name.replace("-", "_").lower()
+    if normalized_recipe_name in {"delayed", "delayed_scaling"}:
+        available, delayed_scaling = safe_import_from(
+            "transformer_engine.common.recipe",
+            "DelayedScaling",
+            msg="model.transformer_engine_fp8=true requires Transformer Engine DelayedScaling",
+        )
+        if not available:
+            raise ImportError("model.transformer_engine_fp8=true requires Transformer Engine DelayedScaling")
+        return delayed_scaling(amax_history_len=amax_history_len, amax_compute_algo=amax_compute_algo)
+
+    if normalized_recipe_name in {"current", "current_scaling"}:
+        available, current_scaling = safe_import_from(
+            "transformer_engine.common.recipe",
+            "Float8CurrentScaling",
+            msg="model.transformer_engine_fp8_recipe=current requires Transformer Engine Float8CurrentScaling",
+        )
+        if not available:
+            raise ImportError("model.transformer_engine_fp8_recipe=current requires Float8CurrentScaling")
+        return current_scaling()
+
+    if normalized_recipe_name in {"mxfp8", "mx", "mx_fp8"}:
+        available, mxfp8_block_scaling = safe_import_from(
+            "transformer_engine.common.recipe",
+            "MXFP8BlockScaling",
+            msg="model.transformer_engine_fp8_recipe=mxfp8 requires Transformer Engine MXFP8BlockScaling",
+        )
+        if not available:
+            raise ImportError("model.transformer_engine_fp8_recipe=mxfp8 requires MXFP8BlockScaling")
+        return mxfp8_block_scaling()
+
+    raise ValueError(
+        f"model.transformer_engine_fp8_recipe must be one of 'delayed', 'current', or 'mxfp8', got {recipe_name!r}"
+    )
+
+
+def _resolve_transformer_engine_autocast() -> Any:
+    """Resolve Transformer Engine's quantization autocast context manager."""
+    available, te_autocast = safe_import_from(
+        "transformer_engine.pytorch.quantization",
+        "autocast",
+        msg="model.transformer_engine_fp8=true requires transformer_engine.pytorch.quantization.autocast",
+    )
+    if not available:
+        raise ImportError("model.transformer_engine_fp8=true requires transformer_engine.pytorch.quantization.autocast")
+    return te_autocast
+
+
 def build_model_and_optimizer(
     *,
     model_id: str,
@@ -199,6 +359,10 @@ def build_model_and_optimizer(
     ddp_cfg: Optional[Dict[str, Any]] = None,
     attention_backend: Optional[str] = None,
     optimizer_cfg: Optional[Dict[str, Any]] = None,
+    transformer_engine_linear: bool = False,
+    transformer_engine_fp8_safe_only: bool = False,
+    fuse_qkv_projections: bool = False,
+    compact_fused_qkv_projections: bool = False,
     pipeline_spec: Optional[Dict[str, Any]] = None,
     peft_cfg=None,
     model_type=None,
@@ -217,6 +381,10 @@ def build_model_and_optimizer(
         ddp_cfg: DDP configuration dict. Mutually exclusive with fsdp_cfg.
         attention_backend: Optional attention backend override.
         optimizer_cfg: Optional optimizer configuration.
+        transformer_engine_linear: Whether to replace transformer torch.nn.Linear modules with Transformer Engine Linear.
+        transformer_engine_fp8_safe_only: Whether to skip TE conversion for known FP8-incompatible modules.
+        fuse_qkv_projections: Whether to call Diffusers QKV projection fusion on the transformer before FSDP.
+        compact_fused_qkv_projections: Whether to remove original projection modules after QKV fusion.
         pipeline_spec: Pipeline specification for pretraining (from_config).
             Required when finetune_mode is False. Should contain:
             - transformer_cls: str (e.g., "WanTransformer3DModel", "FluxTransformer2DModel")
@@ -328,6 +496,10 @@ def build_model_and_optimizer(
             low_cpu_mem_usage=True,
             peft_cfg=peft_cfg,
             model_type=model_type,
+            transformer_engine_linear=transformer_engine_linear,
+            transformer_engine_fp8_safe_only=transformer_engine_fp8_safe_only,
+            fuse_qkv_projections=fuse_qkv_projections,
+            compact_fused_qkv_projections=compact_fused_qkv_projections,
         )
     else:
         # Pretraining: initialize with random weights using pipeline_spec
@@ -347,12 +519,17 @@ def build_model_and_optimizer(
             device=device,
             parallel_scheme=parallel_scheme,
             components_to_load=["transformer"],
+            transformer_engine_linear=transformer_engine_linear,
+            transformer_engine_fp8_safe_only=transformer_engine_fp8_safe_only,
+            fuse_qkv_projections=fuse_qkv_projections,
+            compact_fused_qkv_projections=compact_fused_qkv_projections,
         )
     fsdp2_manager = created_managers["transformer"]
     transformer_module = pipe.transformer
+    transformer_module_for_attrs = getattr(transformer_module, "module", transformer_module)
     if attention_backend is not None:
         logging.info(f"[INFO] Setting attention backend to {attention_backend}")
-        transformer_module.set_attention_backend(attention_backend)
+        transformer_module_for_attrs.set_attention_backend(attention_backend)
 
     if lora_enabled:
         # Collect lora_params AFTER FSDP2 wrapping from the live wrapped module.
@@ -509,6 +686,21 @@ def setup(self):
 
         self.model_id = self.cfg.get("model.pretrained_model_name_or_path")
         self.attention_backend = self.cfg.get("model.attention_backend")
+        self.transformer_engine_linear = bool(self.cfg.get("model.transformer_engine_linear", False))
+        self.transformer_engine_fp8 = bool(self.cfg.get("model.transformer_engine_fp8", False))
+        self.transformer_engine_fp8_recipe_name = str(self.cfg.get("model.transformer_engine_fp8_recipe", "delayed"))
+        self.transformer_engine_fp8_amax_history_len = int(
+            self.cfg.get("model.transformer_engine_fp8_amax_history_len", 1024)
+        )
+        self.transformer_engine_fp8_amax_compute_algo = str(
+            self.cfg.get("model.transformer_engine_fp8_amax_compute_algo", "max")
+        )
+        self.fuse_qkv_projections = bool(self.cfg.get("model.fuse_qkv_projections", False))
+        self.compact_fused_qkv_projections = bool(self.cfg.get("model.compact_fused_qkv_projections", False))
+        if self.transformer_engine_fp8:
+            self.transformer_engine_linear = True
+        if self.compact_fused_qkv_projections and not self.fuse_qkv_projections:
+            raise ValueError("model.compact_fused_qkv_projections=true requires model.fuse_qkv_projections=true")
         self.learning_rate = self.cfg.get("optim.learning_rate", 5e-6)
         self.clip_grad_max_norm = float(self.cfg.get("optim.clip_grad", 1.0))
         self.model_dtype, self.compute_dtype = _resolve_model_dtypes(self.cfg)
@@ -527,6 +719,17 @@ def setup(self):
         self.local_world_size = max(self.local_world_size, 1)
         self.num_nodes = max(1, self.world_size // self.local_world_size)
         self.node_rank = dist.get_rank() // self.local_world_size if dist.is_initialized() else 0
+        self._te_fp8_autocast = None
+        self._te_fp8_recipe = None
+        self._te_fp8_group = None
+        if self.transformer_engine_fp8:
+            self._te_fp8_autocast = _resolve_transformer_engine_autocast()
+            self._te_fp8_recipe = _build_transformer_engine_fp8_recipe(
+                self.transformer_engine_fp8_recipe_name,
+                amax_history_len=self.transformer_engine_fp8_amax_history_len,
+                amax_compute_algo=self.transformer_engine_fp8_amax_compute_algo,
+            )
+            self._te_fp8_group = dist.group.WORLD if dist.is_initialized() else None
 
         logging.info("[INFO] Diffusion Trainer with Flow Matching")
         logging.info(
@@ -534,6 +737,16 @@ def setup(self):
         )
         logging.info(f"[INFO] Node rank: {self.node_rank}, Local rank: {self.local_rank}")
         logging.info(f"[INFO] Learning rate: {self.learning_rate}")
+        logging.info("[INFO] Transformer Engine Linear: %s", self.transformer_engine_linear)
+        logging.info(
+            "[INFO] Transformer Engine FP8: %s (recipe=%s, amax_history_len=%s, amax_compute_algo=%s)",
+            self.transformer_engine_fp8,
+            self.transformer_engine_fp8_recipe_name,
+            self.transformer_engine_fp8_amax_history_len,
+            self.transformer_engine_fp8_amax_compute_algo,
+        )
+        logging.info("[INFO] Fuse QKV projections: %s", self.fuse_qkv_projections)
+        logging.info("[INFO] Compact fused QKV projections: %s", self.compact_fused_qkv_projections)
         logging.info("[INFO] Precision: model_dtype=%s, compute_dtype=%s", self.model_dtype, self.compute_dtype)
         logging.info(
             "[INFO] Performance config: check_loss=%s, grad_clip_foreach=%s",
@@ -626,6 +839,10 @@ def setup(self):
             fsdp_cfg=fsdp_cfg,
             ddp_cfg=ddp_cfg,
             optimizer_cfg=self.cfg.get("optim.optimizer", {}),
+            transformer_engine_linear=self.transformer_engine_linear,
+            transformer_engine_fp8_safe_only=self.transformer_engine_fp8,
+            fuse_qkv_projections=self.fuse_qkv_projections,
+            compact_fused_qkv_projections=self.compact_fused_qkv_projections,
             attention_backend=self.attention_backend,
             pipeline_spec=pipeline_spec,
             peft_cfg=self.peft_cfg,
@@ -769,6 +986,16 @@ def setup(self):
         if dist.is_initialized():
             dist.barrier()
 
+    def _transformer_engine_fp8_context(self) -> Any:
+        """Return the per-forward Transformer Engine FP8 context."""
+        if not self.transformer_engine_fp8:
+            return nullcontext()
+        return self._te_fp8_autocast(
+            enabled=True,
+            recipe=self._te_fp8_recipe,
+            amax_reduction_group=self._te_fp8_group,
+        )
+
     def run_train_validation_loop(self):
         logging.info("[INFO] Starting T2V training with Flow Matching")
         logging.info(f"[INFO] Global Batch size: {self.global_batch_size}; Local Batch size: {self.local_batch_size}")
@@ -809,15 +1036,16 @@ def run_train_validation_loop(self):
                         prepare_for_final_backward([self.model], pp_enabled=False)
 
                     try:
-                        _, average_weighted_loss, _, _ = self.flow_matching_pipeline.step(
-                            model=self.model,
-                            batch=micro_batch,
-                            device=self.device,
-                            dtype=self.compute_dtype,
-                            global_step=global_step,
-                            collect_metrics=False,
-                            check_loss=self.check_loss,
-                        )
+                        with self._transformer_engine_fp8_context():
+                            _, average_weighted_loss, _, _ = self.flow_matching_pipeline.step(
+                                model=self.model,
+                                batch=micro_batch,
+                                device=self.device,
+                                dtype=self.compute_dtype,
+                                global_step=global_step,
+                                collect_metrics=False,
+                                check_loss=self.check_loss,
+                            )
                     except Exception as exc:
                         logging.info(f"[ERROR] Training step failed at epoch {epoch}, step {num_steps}: {exc}")
                         video_shape = micro_batch.get("video_latents", torch.tensor([])).shape
diff --git a/tests/unit_tests/_diffusers/test_auto_diffusion_pipeline.py b/tests/unit_tests/_diffusers/test_auto_diffusion_pipeline.py
index 565ea419b1..1c39f047ae 100644
--- a/tests/unit_tests/_diffusers/test_auto_diffusion_pipeline.py
+++ b/tests/unit_tests/_diffusers/test_auto_diffusion_pipeline.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import logging
+from types import SimpleNamespace
 from unittest.mock import MagicMock, Mock, patch
 
 import pytest
@@ -172,6 +173,145 @@ def test_ensure_params_trainable_returns_correct_count():
     assert count == expected
 
 
+# =============================================================================
+# FP8/linear helper tests
+# =============================================================================
+
+
+def test_is_fp8_training_safe_linear_accepts_multiple_of_16_shapes():
+    from nemo_automodel._diffusers.auto_diffusion_pipeline import _is_fp8_training_safe_linear
+
+    assert _is_fp8_training_safe_linear("block.attn.to_q", torch.nn.Linear(32, 64))
+
+
+@pytest.mark.parametrize(
+    ("name", "linear"),
+    [
+        ("block.attn.to_q", torch.nn.Linear(30, 64)),
+        ("time_text_embed.proj", torch.nn.Linear(32, 64)),
+        ("norm_out.linear", torch.nn.Linear(32, 64)),
+        ("block.norm.linear", torch.nn.Linear(32, 64)),
+    ],
+)
+def test_is_fp8_training_safe_linear_rejects_known_unsafe_cases(name, linear):
+    from nemo_automodel._diffusers.auto_diffusion_pipeline import _is_fp8_training_safe_linear
+
+    assert not _is_fp8_training_safe_linear(name, linear)
+
+
+def test_replace_linear_with_transformer_engine_copies_weights_and_skips_unsafe_modules():
+    from nemo_automodel._diffusers.auto_diffusion_pipeline import _replace_linear_with_transformer_engine
+
+    class FakeTransformerEngineLinear(torch.nn.Linear):
+        def __init__(self, in_features, out_features, bias=True, device=None, params_dtype=None):
+            super().__init__(in_features, out_features, bias=bias, device=device, dtype=params_dtype)
+
+    class MixedLinearModule(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.safe = torch.nn.Linear(16, 32)
+            self.unsafe = torch.nn.Linear(15, 32, bias=False)
+            self.nested = torch.nn.Sequential(torch.nn.Linear(16, 16, bias=False))
+
+    module = MixedLinearModule()
+    module.eval()
+    safe_weight = module.safe.weight.detach().clone()
+    nested_weight = module.nested[0].weight.detach().clone()
+    module.safe.weight.requires_grad_(False)
+
+    fake_transformer_engine = SimpleNamespace(pytorch=SimpleNamespace(Linear=FakeTransformerEngineLinear))
+    with patch(f"{MODULE_PATH}.safe_import_te", return_value=(True, fake_transformer_engine)):
+        converted = _replace_linear_with_transformer_engine(module, "transformer", fp8_safe_only=True)
+
+    assert converted == 2
+    assert isinstance(module.safe, FakeTransformerEngineLinear)
+    assert isinstance(module.nested[0], FakeTransformerEngineLinear)
+    assert isinstance(module.unsafe, torch.nn.Linear)
+    assert not isinstance(module.unsafe, FakeTransformerEngineLinear)
+    assert not module.safe.training
+    assert not module.safe.weight.requires_grad
+    assert torch.allclose(module.safe.weight, safe_weight)
+    assert torch.allclose(module.nested[0].weight, nested_weight)
+
+
+def test_replace_linear_with_transformer_engine_requires_optional_dependency():
+    from nemo_automodel._diffusers.auto_diffusion_pipeline import _replace_linear_with_transformer_engine
+
+    with patch(f"{MODULE_PATH}.safe_import_te", return_value=(False, None)):
+        with pytest.raises(ImportError, match="transformer_engine.pytorch"):
+            _replace_linear_with_transformer_engine(DummyModule(), "transformer")
+
+
+def test_compact_fused_qkv_projection_modules_removes_only_fused_original_projections():
+    from nemo_automodel._diffusers.auto_diffusion_pipeline import _compact_fused_qkv_projection_modules
+
+    class AttentionModule(torch.nn.Module):
+        def __init__(self, *, fused_projections):
+            super().__init__()
+            self.fused_projections = fused_projections
+            self.to_qkv = torch.nn.Linear(2, 6, bias=False)
+            self.to_q = torch.nn.Linear(2, 2, bias=False)
+            self.to_k = torch.nn.Linear(2, 2, bias=False)
+            self.to_v = torch.nn.Linear(2, 2, bias=False)
+            self.to_added_qkv = torch.nn.Linear(2, 6, bias=False)
+            self.add_q_proj = torch.nn.Linear(2, 2, bias=False)
+            self.add_k_proj = torch.nn.Linear(2, 2, bias=False)
+            self.add_v_proj = torch.nn.Linear(2, 2, bias=False)
+
+    module = torch.nn.Module()
+    module.fused = AttentionModule(fused_projections=True)
+    module.unfused = AttentionModule(fused_projections=False)
+    expected_removed_parameters = sum(
+        parameter.numel()
+        for projection_name in ("to_q", "to_k", "to_v", "add_q_proj", "add_k_proj", "add_v_proj")
+        for parameter in getattr(module.fused, projection_name).parameters()
+    )
+
+    removed_modules, removed_parameters = _compact_fused_qkv_projection_modules(module, "transformer")
+
+    assert removed_modules == 6
+    assert removed_parameters == expected_removed_parameters
+    for projection_name in ("to_q", "to_k", "to_v", "add_q_proj", "add_k_proj", "add_v_proj"):
+        assert not hasattr(module.fused, projection_name)
+        assert hasattr(module.unfused, projection_name)
+    assert hasattr(module.fused, "to_qkv")
+    assert hasattr(module.fused, "to_added_qkv")
+
+
+def test_fuse_transformer_qkv_projections_calls_hook_and_counts_fused_modules():
+    from nemo_automodel._diffusers.auto_diffusion_pipeline import _fuse_transformer_qkv_projections
+
+    class AttentionModule(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.fused_projections = False
+
+    class TransformerModule(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.attn1 = AttentionModule()
+            self.attn2 = AttentionModule()
+
+        def fuse_qkv_projections(self):
+            self.attn1.fused_projections = True
+            self.attn2.fused_projections = True
+
+    transformer = TransformerModule()
+
+    fused = _fuse_transformer_qkv_projections(transformer, "transformer")
+
+    assert fused == 2
+    assert transformer.attn1.fused_projections
+    assert transformer.attn2.fused_projections
+
+
+def test_fuse_transformer_qkv_projections_requires_diffusers_hook():
+    from nemo_automodel._diffusers.auto_diffusion_pipeline import _fuse_transformer_qkv_projections
+
+    with pytest.raises(AttributeError, match="fuse_qkv_projections"):
+        _fuse_transformer_qkv_projections(DummyModule(), "transformer")
+
+
 # =============================================================================
 # PipelineSpec tests
 # =============================================================================
@@ -571,6 +711,51 @@ def test_from_pretrained_components_to_load_filters_modules():
     assert mock_move.call_args[0][0] is unet
 
 
+def test_from_pretrained_applies_transformer_qkv_and_te_options():
+    from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline
+
+    transformer = DummyModule()
+    dummy_pipe = DummyPipeline({"transformer": transformer, "text_encoder": DummyModule()})
+    mock_diffusion_pipeline = MagicMock()
+    mock_diffusion_pipeline.from_pretrained.return_value = dummy_pipe
+
+    with (
+        patch(f"{MODULE_PATH}.DIFFUSERS_AVAILABLE", True),
+        patch(f"{MODULE_PATH}.DiffusionPipeline", mock_diffusion_pipeline),
+        patch(f"{MODULE_PATH}._fuse_transformer_qkv_projections") as mock_fuse,
+        patch(f"{MODULE_PATH}._replace_linear_with_transformer_engine") as mock_replace_linear,
+    ):
+        NeMoAutoDiffusionPipeline.from_pretrained(
+            "dummy",
+            move_to_device=False,
+            fuse_qkv_projections=True,
+            compact_fused_qkv_projections=True,
+            transformer_engine_linear=True,
+            transformer_engine_fp8_safe_only=True,
+        )
+
+    mock_fuse.assert_called_once_with(transformer, "transformer", compact=True)
+    mock_replace_linear.assert_called_once_with(transformer, "transformer", fp8_safe_only=True)
+
+
+def test_from_pretrained_rejects_compact_qkv_without_fusion():
+    from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline
+
+    mock_diffusion_pipeline = MagicMock()
+    mock_diffusion_pipeline.from_pretrained.return_value = DummyPipeline({"transformer": DummyModule()})
+
+    with (
+        patch(f"{MODULE_PATH}.DIFFUSERS_AVAILABLE", True),
+        patch(f"{MODULE_PATH}.DiffusionPipeline", mock_diffusion_pipeline),
+    ):
+        with pytest.raises(ValueError, match="compact_fused_qkv_projections=true"):
+            NeMoAutoDiffusionPipeline.from_pretrained(
+                "dummy",
+                move_to_device=False,
+                compact_fused_qkv_projections=True,
+            )
+
+
 # =============================================================================
 # from_config tests
 # =============================================================================
@@ -619,6 +804,26 @@ def test_from_config_raises_without_transformer_cls():
             NeMoAutoDiffusionPipeline.from_config("dummy", pipeline_spec={})
 
 
+def test_from_config_rejects_compact_qkv_without_fusion():
+    from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline
+
+    mock_transformer_cls = MagicMock()
+    mock_transformer_cls.load_config.return_value = {}
+    mock_transformer_cls.from_config.return_value = DummyModule()
+
+    with (
+        patch(f"{MODULE_PATH}.DIFFUSERS_AVAILABLE", True),
+        patch(f"{MODULE_PATH}._import_diffusers_class", return_value=mock_transformer_cls),
+    ):
+        with pytest.raises(ValueError, match="compact_fused_qkv_projections=true"):
+            NeMoAutoDiffusionPipeline.from_config(
+                "model-id",
+                pipeline_spec={"transformer_cls": "FakeTransformer"},
+                move_to_device=False,
+                compact_fused_qkv_projections=True,
+            )
+
+
 def test_from_config_full_pipeline_mode():
     from nemo_automodel._diffusers.auto_diffusion_pipeline import NeMoAutoDiffusionPipeline
 
diff --git a/tests/unit_tests/recipes/test_diffusion_train_metrics.py b/tests/unit_tests/recipes/test_diffusion_train_metrics.py
index 353e2d1e66..5cb536a37a 100644
--- a/tests/unit_tests/recipes/test_diffusion_train_metrics.py
+++ b/tests/unit_tests/recipes/test_diffusion_train_metrics.py
@@ -93,6 +93,207 @@ def test_calculate_throughput_metrics_clamps_invalid_inputs():
     assert metrics["log_window_samples"] == pytest.approx(0.0)
 
 
+@pytest.mark.parametrize(
+    ("value", "expected"),
+    [
+        ("bf16", torch.bfloat16),
+        ("torch.float32", torch.float32),
+        ("unknown", "unknown"),
+        (0.125, 0.125),
+    ],
+)
+def test_normalize_optimizer_value_converts_dtype_aliases(value, expected):
+    assert diffusion_train._normalize_optimizer_value(value) == expected
+
+
+def test_resolve_optimizer_class_handles_default_callable_and_imported_targets(monkeypatch):
+    class CustomOptimizer:
+        pass
+
+    def fake_safe_import_from(module_name, symbol_name, msg):
+        assert module_name == "custom.optim"
+        assert symbol_name == "CustomOptimizer"
+        assert "custom.optim.CustomOptimizer" in msg
+        return True, CustomOptimizer
+
+    monkeypatch.setattr(diffusion_train, "safe_import_from", fake_safe_import_from)
+
+    assert diffusion_train._resolve_optimizer_class("torch.optim.AdamW") is torch.optim.AdamW
+    assert diffusion_train._resolve_optimizer_class(CustomOptimizer) is CustomOptimizer
+    assert diffusion_train._resolve_optimizer_class("custom.optim.CustomOptimizer") is CustomOptimizer
+
+
+@pytest.mark.parametrize("target", ["AdamW", object()])
+def test_resolve_optimizer_class_rejects_invalid_targets(target):
+    with pytest.raises(ValueError, match="Optimizer target must be"):
+        diffusion_train._resolve_optimizer_class(target)
+
+
+def test_resolve_optimizer_class_raises_when_import_fails(monkeypatch):
+    monkeypatch.setattr(diffusion_train, "safe_import_from", lambda *_args, **_kwargs: (False, None))
+
+    with pytest.raises(ImportError, match="could not be imported"):
+        diffusion_train._resolve_optimizer_class("missing.optim.CustomOptimizer")
+
+
+def test_filter_optimizer_kwargs_keeps_only_supported_parameters():
+    class OptimizerWithoutKwargs:
+        def __init__(self, params, lr=0.1, *, beta=0.9):
+            self.params = params
+            self.lr = lr
+            self.beta = beta
+
+    optimizer_kwargs = {"lr": 0.01, "beta": 0.95, "weight_decay": 0.1}
+
+    assert diffusion_train._filter_optimizer_kwargs(
+        "custom.OptimizerWithoutKwargs", OptimizerWithoutKwargs, optimizer_kwargs
+    ) == {"lr": 0.01, "beta": 0.95}
+
+
+def test_filter_optimizer_kwargs_passes_all_kwargs_when_target_accepts_var_kwargs():
+    class OptimizerWithKwargs:
+        def __init__(self, params, **kwargs):
+            self.params = params
+            self.kwargs = kwargs
+
+    optimizer_kwargs = {"lr": 0.01, "weight_decay": 0.1}
+
+    assert (
+        diffusion_train._filter_optimizer_kwargs("custom.OptimizerWithKwargs", OptimizerWithKwargs, optimizer_kwargs)
+        is optimizer_kwargs
+    )
+
+
+def test_filter_optimizer_kwargs_passes_all_kwargs_when_signature_cannot_be_inspected(monkeypatch):
+    optimizer_kwargs = {"lr": 0.01, "weight_decay": 0.1}
+
+    def raise_value_error(_target):
+        raise ValueError("no signature")
+
+    monkeypatch.setattr(diffusion_train.inspect, "signature", raise_value_error)
+
+    assert (
+        diffusion_train._filter_optimizer_kwargs("custom.NoSignature", object(), optimizer_kwargs) is optimizer_kwargs
+    )
+
+
+def test_build_transformer_engine_fp8_recipe_dispatches_recipe_names(monkeypatch):
+    class DelayedScaling:
+        def __init__(self, *, amax_history_len, amax_compute_algo):
+            self.amax_history_len = amax_history_len
+            self.amax_compute_algo = amax_compute_algo
+
+    class Float8CurrentScaling:
+        pass
+
+    class MXFP8BlockScaling:
+        pass
+
+    recipe_classes = {
+        "DelayedScaling": DelayedScaling,
+        "Float8CurrentScaling": Float8CurrentScaling,
+        "MXFP8BlockScaling": MXFP8BlockScaling,
+    }
+
+    def fake_safe_import_from(module_name, symbol_name, msg):
+        assert module_name == "transformer_engine.common.recipe"
+        assert "Transformer Engine" in msg
+        return True, recipe_classes[symbol_name]
+
+    monkeypatch.setattr(diffusion_train, "safe_import_from", fake_safe_import_from)
+
+    delayed = diffusion_train._build_transformer_engine_fp8_recipe(
+        "delayed-scaling",
+        amax_history_len=32,
+        amax_compute_algo="max",
+    )
+    current = diffusion_train._build_transformer_engine_fp8_recipe(
+        "current",
+        amax_history_len=32,
+        amax_compute_algo="max",
+    )
+    mxfp8 = diffusion_train._build_transformer_engine_fp8_recipe(
+        "mx",
+        amax_history_len=32,
+        amax_compute_algo="max",
+    )
+
+    assert isinstance(delayed, DelayedScaling)
+    assert delayed.amax_history_len == 32
+    assert delayed.amax_compute_algo == "max"
+    assert isinstance(current, Float8CurrentScaling)
+    assert isinstance(mxfp8, MXFP8BlockScaling)
+
+
+def test_build_transformer_engine_fp8_recipe_rejects_unknown_recipe():
+    with pytest.raises(ValueError, match="must be one of"):
+        diffusion_train._build_transformer_engine_fp8_recipe(
+            "unknown",
+            amax_history_len=32,
+            amax_compute_algo="max",
+        )
+
+
+def test_build_transformer_engine_fp8_recipe_requires_optional_dependency(monkeypatch):
+    monkeypatch.setattr(diffusion_train, "safe_import_from", lambda *_args, **_kwargs: (False, None))
+
+    with pytest.raises(ImportError, match="DelayedScaling"):
+        diffusion_train._build_transformer_engine_fp8_recipe(
+            "delayed",
+            amax_history_len=32,
+            amax_compute_algo="max",
+        )
+
+
+def test_resolve_transformer_engine_autocast_returns_imported_context_manager(monkeypatch):
+    autocast = object()
+
+    def fake_safe_import_from(module_name, symbol_name, msg):
+        assert module_name == "transformer_engine.pytorch.quantization"
+        assert symbol_name == "autocast"
+        assert "autocast" in msg
+        return True, autocast
+
+    monkeypatch.setattr(diffusion_train, "safe_import_from", fake_safe_import_from)
+
+    assert diffusion_train._resolve_transformer_engine_autocast() is autocast
+
+
+def test_resolve_transformer_engine_autocast_requires_optional_dependency(monkeypatch):
+    monkeypatch.setattr(diffusion_train, "safe_import_from", lambda *_args, **_kwargs: (False, None))
+
+    with pytest.raises(ImportError, match="autocast"):
+        diffusion_train._resolve_transformer_engine_autocast()
+
+
+def test_transformer_engine_fp8_context_returns_nullcontext_when_disabled():
+    recipe = object.__new__(TrainDiffusionRecipe)
+    recipe.transformer_engine_fp8 = False
+
+    with recipe._transformer_engine_fp8_context():
+        pass
+
+
+def test_transformer_engine_fp8_context_calls_transformer_engine_autocast():
+    recipe = object.__new__(TrainDiffusionRecipe)
+    context = MagicMock()
+    context.__enter__.return_value = None
+    context.__exit__.return_value = None
+    recipe.transformer_engine_fp8 = True
+    recipe._te_fp8_autocast = MagicMock(return_value=context)
+    recipe._te_fp8_recipe = "recipe"
+    recipe._te_fp8_group = "group"
+
+    with recipe._transformer_engine_fp8_context():
+        pass
+
+    recipe._te_fp8_autocast.assert_called_once_with(
+        enabled=True,
+        recipe="recipe",
+        amax_reduction_group="group",
+    )
+
+
 class _TinyTransformer(nn.Module):
     def __init__(self):
         super().__init__()
@@ -139,6 +340,10 @@ def fake_from_pretrained(model_id, **kwargs):
             "fsdp2_forward_prefetch_depth": 3,
         },
         attention_backend="flash",
+        transformer_engine_linear=True,
+        transformer_engine_fp8_safe_only=True,
+        fuse_qkv_projections=True,
+        compact_fused_qkv_projections=True,
         optimizer_cfg={
             "weight_decay": 0.25,
             "betas": [0.8, 0.95],
@@ -159,6 +364,10 @@ def fake_from_pretrained(model_id, **kwargs):
     assert manager_args["enable_fsdp2_prefetch"] is False
     assert manager_args["fsdp2_backward_prefetch_depth"] == 4
     assert manager_args["fsdp2_forward_prefetch_depth"] == 3
+    assert calls["transformer_engine_linear"] is True
+    assert calls["transformer_engine_fp8_safe_only"] is True
+    assert calls["fuse_qkv_projections"] is True
+    assert calls["compact_fused_qkv_projections"] is True
     assert pipe.transformer.attention_backend == "flash"
     assert device_mesh == "mesh"
     assert optimizer.defaults["lr"] == pytest.approx(0.125)
@@ -332,6 +541,7 @@ def fake_tqdm(iterable, desc):
     recipe.check_loss = True
     recipe.clip_grad_max_norm = 0.5
     recipe.grad_clip_foreach = False
+    recipe.transformer_engine_fp8 = False
     recipe.peft_cfg = None
     recipe.log_every = 1
     recipe._elapsed_seconds_since = MagicMock(return_value=(2.0, 10.0))