feat(maxtext): Maxtext & MaxDiffusion refactor

examples/maxdiffusion/configs/MI300X/flux_dev-pretrain.yaml

@@ -0,0 +1,51 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_dev-pretrain}
+workspace: ./output
+
+# Per-config environment (arch-agnostic; see the MI355X flux config). The
+# gfx950-only RCCL_WARP_SPEED_AUTO=0 workaround is applied by run.sh and is a
+# no-op on MI300X (gfx942). NOTE: as on MI355X, Shardy is left at the image
+# default (the default maxdiffusion image forces it on for cudnn_flash_te); set
+# JAX_USE_SHARDY_PARTITIONER here only if you switch to an older image.
+env:
+  KERAS_BACKEND: "jax"
+  JAX_SPMD_MODE: "allow_all"
+  JAX_TRACEBACK_FILTERING: "off"
+  JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: "all"
+  TOKENIZERS_PARALLELISM: "1"
+  SKIP_GCS: "1"
+  XLA_PYTHON_CLIENT_MEM_FRACTION: "0.95"
+  TF_CUDNN_WORKSPACE_LIMIT_IN_MB: "71296"
+  XLA_FLAGS: "--xla_gpu_enable_latency_hiding_scheduler=true --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=5 --xla_gpu_enable_reduce_scatter_combine_by_dim=false --xla_gpu_enable_all_gather_combine_by_dim=false --xla_gpu_all_gather_combine_threshold_bytes=134217728 --xla_gpu_reduce_scatter_combine_threshold_bytes=134217728 --xla_gpu_enable_command_buffer=''"
+  NVTE_FUSED_ATTN: "1"
+  NVTE_FUSED_ATTN_CK: "1"
+  NVTE_FUSED_ATTN_AOTRITON: "0"
+  NVTE_CK_USES_BWD_V3: "1"
+  NVTE_CK_USES_FWD_V3: "1"
+  NVTE_CK_IS_V3_ATOMIC_FP32: "0"
+  NVTE_CK_HOW_V3_BF16_CVT: "2"
+  NVTE_ALLOW_NONDETERMINISTIC_ALGO: "1"
+
+modules:
+  pre_trainer:
+    framework: maxdiffusion
+    config: pre_trainer.yaml
+
+    model: flux_dev.yaml
+    maxdiffusion_entrypoint: flux
+
+    overrides:
+      run_name: "flux_dev_pretrain"
+      output_dir: "./output/flux_dev-pretrain"
+      base_output_directory: "./output/flux_dev-pretrain"
+      max_train_steps: ${MAX_STEPS:20}
+      dataset_type: "synthetic"
+      hardware: "gpu"
+      checkpoint_every: -1
+      enable_profiler: false
+      # MI300X has 192 GB HBM (vs 288 GB on MI355X); the MI355X config uses
+      # per_device_batch_size 14. Lower here if FLUX.1-dev OOMs on MI300X.
+      per_device_batch_size: 8
+      ici_data_parallelism: 1
+      ici_tensor_parallelism: 1

examples/maxdiffusion/configs/MI300X/wan2.1_1.3b-pretrain.yaml

@@ -0,0 +1,52 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:wan2.1_1.3b-pretrain}
+workspace: ./output
+
+# Per-config environment (ported from the retired MAD jax-maxdiffusion env
+# scripts). TrainRuntime applies this top-level `env:` before JAX/XLA init.
+# These flags are arch-agnostic; the gfx950-only RCCL_WARP_SPEED_AUTO=0
+# workaround is applied by run.sh and is a no-op on MI300X (gfx942).
+env:
+  KERAS_BACKEND: "jax"
+  JAX_SPMD_MODE: "allow_all"
+  JAX_TRACEBACK_FILTERING: "off"
+  JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: "all"
+  TOKENIZERS_PARALLELISM: "1"
+  SKIP_GCS: "1"
+  XLA_PYTHON_CLIENT_MEM_FRACTION: "0.95"
+  TF_CUDNN_WORKSPACE_LIMIT_IN_MB: "71296"
+  XLA_FLAGS: "--xla_gpu_enable_latency_hiding_scheduler=true --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=5 --xla_gpu_enable_reduce_scatter_combine_by_dim=false --xla_gpu_enable_all_gather_combine_by_dim=false --xla_gpu_all_gather_combine_threshold_bytes=134217728 --xla_gpu_reduce_scatter_combine_threshold_bytes=134217728 --xla_gpu_enable_command_buffer=''"
+  NVTE_FUSED_ATTN: "1"
+  NVTE_FUSED_ATTN_CK: "1"
+  NVTE_FUSED_ATTN_AOTRITON: "0"
+  NVTE_CK_USES_BWD_V3: "1"
+  NVTE_CK_USES_FWD_V3: "1"
+  NVTE_CK_IS_V3_ATOMIC_FP32: "0"
+  NVTE_CK_HOW_V3_BF16_CVT: "2"
+  NVTE_ALLOW_NONDETERMINISTIC_ALGO: "1"
+
+modules:
+  pre_trainer:
+    framework: maxdiffusion
+    config: pre_trainer.yaml
+
+    model: wan2.1_1.3b.yaml
+    maxdiffusion_entrypoint: wan
+
+    overrides:
+      run_name: "wan2.1_1.3b_pretrain"
+      output_dir: "./output/wan2.1_1.3b-pretrain"
+      base_output_directory: "./output/wan2.1_1.3b-pretrain"
+      max_train_steps: ${MAX_STEPS:20}
+      dataset_type: "synthetic"
+      hardware: "gpu"
+      checkpoint_every: -1
+      enable_profiler: false
+      # MI300X has 192 GB HBM (vs 288 GB on MI355X); reduce per_device_batch_size
+      # here if a model OOMs relative to the MI355X config.
+      per_device_batch_size: 1
+      ici_fsdp_parallelism: 8
+      ici_data_parallelism: 1
+      ici_context_parallelism: 1
+      ici_tensor_parallelism: 1

examples/maxdiffusion/configs/MI300X/wan2.1_14b-pretrain.yaml

@@ -0,0 +1,51 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:wan2.1_14b-pretrain}
+workspace: ./output
+
+# Per-config environment (arch-agnostic; see the 1.3B MI300X config). The
+# gfx950-only RCCL_WARP_SPEED_AUTO=0 workaround is applied by run.sh and is a
+# no-op on MI300X (gfx942).
+env:
+  KERAS_BACKEND: "jax"
+  JAX_SPMD_MODE: "allow_all"
+  JAX_TRACEBACK_FILTERING: "off"
+  JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: "all"
+  TOKENIZERS_PARALLELISM: "1"
+  SKIP_GCS: "1"
+  XLA_PYTHON_CLIENT_MEM_FRACTION: "0.95"
+  TF_CUDNN_WORKSPACE_LIMIT_IN_MB: "71296"
+  XLA_FLAGS: "--xla_gpu_enable_latency_hiding_scheduler=true --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=5 --xla_gpu_enable_reduce_scatter_combine_by_dim=false --xla_gpu_enable_all_gather_combine_by_dim=false --xla_gpu_all_gather_combine_threshold_bytes=134217728 --xla_gpu_reduce_scatter_combine_threshold_bytes=134217728 --xla_gpu_enable_command_buffer=''"
+  NVTE_FUSED_ATTN: "1"
+  NVTE_FUSED_ATTN_CK: "1"
+  NVTE_FUSED_ATTN_AOTRITON: "0"
+  NVTE_CK_USES_BWD_V3: "1"
+  NVTE_CK_USES_FWD_V3: "1"
+  NVTE_CK_IS_V3_ATOMIC_FP32: "0"
+  NVTE_CK_HOW_V3_BF16_CVT: "2"
+  NVTE_ALLOW_NONDETERMINISTIC_ALGO: "1"
+
+modules:
+  pre_trainer:
+    framework: maxdiffusion
+    config: pre_trainer.yaml
+
+    model: wan2.1_14b.yaml
+    maxdiffusion_entrypoint: wan
+
+    overrides:
+      run_name: "wan2.1_14b_pretrain"
+      output_dir: "./output/wan2.1_14b-pretrain"
+      base_output_directory: "./output/wan2.1_14b-pretrain"
+      max_train_steps: ${MAX_STEPS:20}
+      dataset_type: "synthetic"
+      hardware: "gpu"
+      checkpoint_every: -1
+      enable_profiler: false
+      # MI300X has 192 GB HBM (vs 288 GB on MI355X); reduce per_device_batch_size
+      # here if wan2.1_14b OOMs relative to the MI355X config.
+      per_device_batch_size: 1
+      ici_fsdp_parallelism: 8
+      ici_data_parallelism: 1
+      ici_context_parallelism: 1
+      ici_tensor_parallelism: 1

examples/maxdiffusion/configs/MI355X/flux_dev-pretrain.yaml

@@ -0,0 +1,82 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:flux_dev-pretrain}
+workspace: ./output
+
+# Per-config environment (ported from base_flux_dev_env.sh). NOTE: the retired
+# shell env also set JAX_USE_SHARDY_PARTITIONER=0, but the default maxdiffusion
+# image (docker/jax_maxdiffusion.jax_the_rock_ci_45156b7_20260701) patches
+# attention_flax.py to FORCE Shardy on for cudnn_flash_te (GSPMD aborts there on
+# JAX 0.10 + TE 2.14). We therefore leave Shardy at the image default; set
+# JAX_USE_SHARDY_PARTITIONER here only if you switch to an older image.
+env:
+  KERAS_BACKEND: "jax"
+  JAX_SPMD_MODE: "allow_all"
+  JAX_TRACEBACK_FILTERING: "off"
+  JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: "all"
+  # Persist the XLA compilation cache on the bind-mounted workspace (NFS) instead
+  # of the ephemeral in-container /app/.cache/jax, so a cold compile (which can be
+  # tens of minutes with autotune_level=5) is paid once and reused across runs
+  # rather than every run recompiling and risking a launcher timeout. /myworkspace
+  # is the container bind-mount of the repo root.
+  JAX_COMPILATION_CACHE_DIR: "/myworkspace/.jax_cache/flux_dev"
+  MIOPEN_CUSTOM_CACHE_DIR: "/myworkspace/.jax_cache/flux_dev/miopen"
+  TOKENIZERS_PARALLELISM: "1"
+  SKIP_GCS: "1"
+  XLA_PYTHON_CLIENT_MEM_FRACTION: "0.95"
+  TF_CUDNN_WORKSPACE_LIMIT_IN_MB: "71296"
+  XLA_FLAGS: "--xla_gpu_enable_latency_hiding_scheduler=true --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=5 --xla_gpu_enable_reduce_scatter_combine_by_dim=false --xla_gpu_enable_all_gather_combine_by_dim=false --xla_gpu_all_gather_combine_threshold_bytes=134217728 --xla_gpu_reduce_scatter_combine_threshold_bytes=134217728 --xla_gpu_enable_command_buffer=''"
+  # ROCm/HIP + RCCL/NCCL knobs. These MUST match the verified-good
+  # base_flux_dev_env.sh: dropping NCCL_PROTO=Simple / HSA_NO_SCRATCH_RECLAIM=1 /
+  # GPU_MAX_HW_QUEUES / HIP_FORCE_DEV_KERNARG / RCCL_MSCCL_ENABLE caused an RCCL
+  # init hang (GPUs 0%, CPU busy-wait) during model construction.
+  HSA_FORCE_FINE_GRAIN_PCIE: "1"
+  HSA_NO_SCRATCH_RECLAIM: "1"
+  GPU_MAX_HW_QUEUES: "2"
+  HIP_FORCE_DEV_KERNARG: "1"
+  RCCL_MSCCL_ENABLE: "0"
+  NCCL_MAX_NCHANNELS: "112"
+  NCCL_IB_TC: "41"
+  NCCL_IB_SL: "0"
+  NCCL_DEBUG: "WARN"
+  NCCL_PROTO: "Simple"
+  LIBTPU_INIT_ARGS: ""
+  NVTE_FUSED_ATTN: "1"
+  NVTE_FUSED_ATTN_CK: "1"
+  NVTE_FUSED_ATTN_AOTRITON: "0"
+  NVTE_CK_USES_BWD_V3: "1"
+  NVTE_CK_USES_FWD_V3: "1"
+  NVTE_CK_IS_V3_ATOMIC_FP32: "0"
+  NVTE_CK_HOW_V3_BF16_CVT: "2"
+  NVTE_ALLOW_NONDETERMINISTIC_ALGO: "1"
+
+modules:
+  pre_trainer:
+    framework: maxdiffusion
+    config: pre_trainer.yaml
+
+    model: flux_dev.yaml
+    maxdiffusion_entrypoint: flux
+
+    overrides:
+      run_name: "flux_dev_pretrain"
+      output_dir: "./output/flux_dev-pretrain"
+      base_output_directory: "./output/flux_dev-pretrain"
+      max_train_steps: ${MAX_STEPS:20}
+      # Log throughput every step: the perf extractor parses the per-step
+      # "completed step: N, seconds: X, TFLOP/s/device: Y" line. The image
+      # default log_period=100 emits nothing at 20 steps, so perf came back empty
+      # and madengine silently reused a stale CSV.
+      log_period: 1
+      # Bound to run.sh's PERF_METRICS_FILE: the trainer writes per-step JSON
+      # metrics here (max_utils.write_metrics_locally). This is the reliable perf
+      # source; the per-step stdout "completed step:" line is dropped under the
+      # Primus launcher. Empty when unset -> no-op (upstream default).
+      metrics_file: "${PERF_METRICS_FILE:}"
+      dataset_type: "synthetic"
+      hardware: "gpu"
+      checkpoint_every: -1
+      enable_profiler: false
+      per_device_batch_size: 14
+      ici_data_parallelism: 1
+      ici_tensor_parallelism: 1

examples/maxdiffusion/configs/MI355X/wan2.1_1.3b-pretrain.yaml

@@ -0,0 +1,87 @@
+work_group: ${PRIMUS_TEAM:amd}
+user_name: ${PRIMUS_USER:root}
+exp_name: ${PRIMUS_EXP_NAME:wan2.1_1.3b-pretrain}
+workspace: ./output
+
+# Per-config environment (single source of truth in Primus). TrainRuntime applies
+# this top-level `env:` before JAX/XLA init. Ported from the retired MAD
+# scripts/jax-maxdiffusion/env_scripts/base_wan_1.3B_env.sh so the JAX
+# MaxDiffusion tuning travels with the config instead of a sourced shell script.
+env:
+  KERAS_BACKEND: "jax"
+  JAX_SPMD_MODE: "allow_all"
+  JAX_TRACEBACK_FILTERING: "off"
+  JAX_PERSISTENT_CACHE_ENABLE_XLA_CACHES: "all"
+  # Persist the XLA compilation cache on the bind-mounted workspace (NFS) instead
+  # of the ephemeral in-container /app/.cache/jax. The WAN fused-attention graph
+  # (long video sequence + autotune_level=5) can take 30-45 min to compile on a
+  # cold cache; without persistence every run recompiles from scratch and can
+  # exceed the launcher timeout before a single step runs. /myworkspace is the
+  # container bind-mount of the repo root.
+  JAX_COMPILATION_CACHE_DIR: "/myworkspace/.jax_cache/wan2.1_1.3b"
+  MIOPEN_CUSTOM_CACHE_DIR: "/myworkspace/.jax_cache/wan2.1_1.3b/miopen"
+  TOKENIZERS_PARALLELISM: "1"
+  SKIP_GCS: "1"
+  XLA_PYTHON_CLIENT_MEM_FRACTION: "0.95"
+  TF_CUDNN_WORKSPACE_LIMIT_IN_MB: "71296"
+  XLA_FLAGS: "--xla_gpu_enable_latency_hiding_scheduler=true --xla_gpu_enable_cublaslt=True --xla_gpu_autotune_level=5 --xla_gpu_enable_reduce_scatter_combine_by_dim=false --xla_gpu_enable_all_gather_combine_by_dim=false --xla_gpu_all_gather_combine_threshold_bytes=134217728 --xla_gpu_reduce_scatter_combine_threshold_bytes=134217728 --xla_gpu_enable_command_buffer=''"
+  # ROCm/HIP + RCCL/NCCL knobs. These MUST match the verified-good
+  # base_wan_1.3B_env.sh: dropping NCCL_PROTO=Simple / HSA_NO_SCRATCH_RECLAIM=1 /
+  # GPU_MAX_HW_QUEUES / HIP_FORCE_DEV_KERNARG / RCCL_MSCCL_ENABLE caused an RCCL
+  # init hang (GPUs 0%, CPU busy-wait) during model construction.
+  HSA_FORCE_FINE_GRAIN_PCIE: "1"
+  HSA_NO_SCRATCH_RECLAIM: "1"
+  GPU_MAX_HW_QUEUES: "2"
+  HIP_FORCE_DEV_KERNARG: "1"
+  RCCL_MSCCL_ENABLE: "0"
+  NCCL_MAX_NCHANNELS: "112"
+  NCCL_IB_TC: "41"
+  NCCL_IB_SL: "0"
+  NCCL_DEBUG: "WARN"
+  NCCL_PROTO: "Simple"
+  LIBTPU_INIT_ARGS: ""
+  NVTE_FUSED_ATTN: "1"
+  NVTE_FUSED_ATTN_CK: "1"
+  NVTE_FUSED_ATTN_AOTRITON: "0"
+  NVTE_CK_USES_BWD_V3: "1"
+  NVTE_CK_USES_FWD_V3: "1"
+  NVTE_CK_IS_V3_ATOMIC_FP32: "0"
+  NVTE_CK_HOW_V3_BF16_CVT: "2"
+  NVTE_ALLOW_NONDETERMINISTIC_ALGO: "1"
+
+modules:
+  pre_trainer:
+    framework: maxdiffusion
+    config: pre_trainer.yaml
+
+    # MaxDiffusion model preset (the native pyconfig file, hosted in Primus).
+    model: wan2.1_1.3b.yaml
+
+    # Primus-only selector (stripped before the config reaches MaxDiffusion's
+    # pyconfig): picks the src.maxdiffusion.train_wan entrypoint + WanTrainer.
+    maxdiffusion_entrypoint: wan
+
+    overrides:
+      run_name: "wan2.1_1.3b_pretrain"
+      output_dir: "./output/wan2.1_1.3b-pretrain"
+      base_output_directory: "./output/wan2.1_1.3b-pretrain"
+      max_train_steps: ${MAX_STEPS:20}
+      # Log throughput every step: the perf extractor parses the per-step
+      # "completed step: N, seconds: X, TFLOP/s/device: Y" line. The image
+      # default log_period=100 emits nothing at 20 steps, so perf came back empty
+      # and madengine silently reused a stale CSV.
+      log_period: 1
+      # Bound to run.sh's PERF_METRICS_FILE: the trainer writes per-step JSON
+      # metrics here (max_utils.write_metrics_locally). This is the reliable perf
+      # source; the per-step stdout "completed step:" line is dropped under the
+      # Primus launcher. Empty when unset -> no-op (upstream default).
+      metrics_file: "${PERF_METRICS_FILE:}"
+      dataset_type: "synthetic"
+      hardware: "gpu"
+      checkpoint_every: -1
+      enable_profiler: false
+      per_device_batch_size: 1
+      ici_fsdp_parallelism: 8
+      ici_data_parallelism: 1
+      ici_context_parallelism: 1
+      ici_tensor_parallelism: 1