docs: Primus JAX MaxText v26.4 (ROCm#6361)

(cherry picked from commit cd6e604)

Author: peterjunpark

docs/data/how-to/rocm-for-ai/training/jax-maxtext-benchmark-models.yaml

@@ -1,12 +1,12 @@
 dockers:
-  - pull_tag: rocm/jax-training:maxtext-v26.3
-    docker_hub_url: https://hub.docker.com/layers/rocm/jax-training/maxtext-v26.3/images/sha256-548381e496fa1367772964341130cc014ee1a9a5ca715ff4ca3aa92960d162c3
+  - pull_tag: rocm/jax-training:maxtext-v26.4-jax0.9.1-te2.12.0
+    docker_hub_url: https://hub.docker.com/layers/rocm/jax-training/maxtext-v26.4-jax0.9.1-te2.12.0/images/
     components:
-      ROCm: 7.2.1
-      JAX: 0.8.2
-      Python: 3.12
-      Transformer Engine: 2.8.0.dev0+9b312832
-      hipBLASLt: 1.3.0+bfcf25fa18
+      ROCm: 7.14.0a20260526
+      JAX: 0.9.1
+      Python: 3.12.3
+      Transformer Engine: 2.12.0.dev0+635d7c08
+      hipBLASLt: 1.4.0+807283e5
 model_groups:
   - group: Meta Llama
     tag: llama
@@ -58,7 +58,7 @@ model_groups:
         model_repo: Llama-3.1-405B
         precision: bf16
         multinode_config:
-          gfx950: env_scripts/gfx950_llama3_405b.yml
+          gfx950: env_scripts/gfx950_llama3.1_405b.yml
         doc_options: ["multi-node"]
       - model: Llama 3.3 70B
         mad_tag: jax_maxtext_train_llama-3.3-70b

docs/data/how-to/rocm-for-ai/training/previous-versions/jax-maxtext-v26.3-benchmark-models.yaml

@@ -0,0 +1,116 @@
+dockers:
+  - pull_tag: rocm/jax-training:maxtext-v26.3
+    docker_hub_url: https://hub.docker.com/layers/rocm/jax-training/maxtext-v26.3/images/sha256-548381e496fa1367772964341130cc014ee1a9a5ca715ff4ca3aa92960d162c3
+    components:
+      ROCm: 7.2.1
+      JAX: 0.8.2
+      Python: 3.12
+      Transformer Engine: 2.8.0.dev0+9b312832
+      hipBLASLt: 1.3.0+bfcf25fa18
+model_groups:
+  - group: Meta Llama
+    tag: llama
+    models:
+      - model: Llama 2 7B
+        mad_tag: jax_maxtext_train_llama-2-7b
+        model_repo: Llama-2-7B
+        precision: bf16
+        primus_config_name: llama2_7B-pretrain.yaml
+        multinode_config:
+          gfx950: env_scripts/gfx950_llama2_7b.yml
+          gfx942: env_scripts/llama2_7b.yml
+        doc_options: ["single-node", "multi-node"]
+      - model: Llama 2 70B
+        mad_tag: jax_maxtext_train_llama-2-70b
+        model_repo: Llama-2-70B
+        precision: bf16
+        primus_config_name: llama2_70B-pretrain.yaml
+        multinode_config:
+          gfx950: env_scripts/gfx950_llama2_70b.yml
+          gfx942: env_scripts/llama2_70b.yml
+        doc_options: ["single-node", "multi-node"]
+      - model: Llama 3 8B
+        mad_tag: jax_maxtext_train_llama-3-8b
+        primus_config_name: llama3_8B-pretrain.yaml
+        multinode_config:
+          gfx950: env_scripts/gfx950_llama3_8b.yml
+          gfx942: env_scripts/llama3_8b.yml
+        doc_options: ["multi-node"]
+      - model: Llama 3 70B
+        mad_tag: jax_maxtext_train_llama-3-70b
+        primus_config_name: llama3_70B-pretrain.yaml
+        multinode_config:
+          gfx950: env_scripts/gfx950_llama3_70b.yml
+          gfx942: env_scripts/llama3_70b.yml
+        doc_options: ["multi-node"]
+      - model: Llama 3.1 8B
+        mad_tag: jax_maxtext_train_llama-3.1-8b
+        model_repo: Llama-3.1-8B
+        precision: bf16
+        doc_options: ["single-node"]
+      - model: Llama 3.1 70B
+        mad_tag: jax_maxtext_train_llama-3.1-70b
+        model_repo: Llama-3.1-70B
+        precision: bf16
+        doc_options: ["single-node"]
+      - model: Llama 3.1 405B (multi-node)
+        mad_tag: jax_maxtext_train_llama-3.1-405b
+        model_repo: Llama-3.1-405B
+        precision: bf16
+        multinode_config:
+          gfx950: env_scripts/gfx950_llama3_405b.yml
+        doc_options: ["multi-node"]
+      - model: Llama 3.3 70B
+        mad_tag: jax_maxtext_train_llama-3.3-70b
+        model_repo: Llama-3.3-70B
+        precision: bf16
+        primus_config_name: llama3.3_70B-pretrain.yaml
+        multinode_config:
+          gfx950: env_scripts/gfx950_llama3.3_70b.yml
+          gfx942: env_scripts/llama3.3_70b.yml
+        doc_options: ["single-node", "multi-node"]
+  - group: DeepSeek
+    tag: deepseek
+    models:
+      - model: DeepSeek-V2-Lite (16B)
+        mad_tag: jax_maxtext_train_deepseek-v2-lite-16b
+        model_repo: DeepSeek-V2-lite
+        precision: bf16
+        primus_config_name: deepseek_v2_16B-pretrain.yaml
+        multinode_config:
+          gfx950: env_scripts/gfx950_deepseek2_16b.yml
+          gfx942: env_scripts/deepseek2_16b.yml
+        doc_options: ["single-node", "multi-node"]
+  - group: Mistral AI
+    tag: mistral
+    models:
+      - model: Mixtral 8x7B
+        mad_tag: jax_maxtext_train_mixtral-8x7b
+        model_repo: Mixtral-8x7B
+        precision: bf16
+        primus_config_name: mixtral_8x7B-pretrain.yaml
+        multinode_config:
+          gfx950: env_scripts/gfx950_mixtral_8x7b.yml
+          gfx942: env_scripts/llama3_8x7b.yml
+        doc_options: ["single-node", "multi-node"]
+  - group: Qwen
+    tag: qwen
+    models:
+      - model: Qwen 14B
+        mad_tag: jax_maxtext_train_qwen3-14b
+        model_repo: Qwen3-14B
+        precision: bf16
+        primus_config_name: qwen3_14B-pretrain.yaml
+        multinode_config:
+          gfx950: env_scripts/gfx950_qwen3_14b.yml
+          gfx942: env_scripts/qwen3_14b.yml
+        doc_options: ["single-node", "multi-node"]
+      - model: Qwen 30B A3B
+        mad_tag: jax_maxtext_train_qwen3-30b-a3b
+        model_repo: Qwen3-30B-A3B
+        precision: bf16
+        primus_config_name: qwen3_30B_a3b-pretrain.yaml
+        multinode_config:
+          gfx950: env_scripts/gfx950_qwen3_30b_a3b.yml
+          gfx942: env_scripts/qwen3_30b_a3b.yml
+        doc_options: ["single-node", "multi-node"]

docs/how-to/rocm-for-ai/training/benchmark-docker/jax-maxtext.rst

@@ -52,7 +52,7 @@ MaxText with on ROCm provides the following key features to train large language
 
 - NANOO FP8 (for MI300X series GPUs) and FP8 (for MI355X and MI350X) quantization support
 
-.. _amd-maxtext-model-support-v26.3:
+.. _amd-maxtext-model-support-v26.4:
 
 Supported models
 ================
@@ -134,15 +134,15 @@ Use the following command to pull the Docker image from Docker Hub.
 
       docker pull {{ docker.pull_tag }}
 
-.. _amd-maxtext-multi-node-setup-v26.3:
+.. _amd-maxtext-multi-node-setup-v26.4:
 
 Multi-node configuration
 ------------------------
 
 See :doc:`/how-to/rocm-for-ai/system-setup/multi-node-setup` to configure your
 environment for multi-node training.
 
-.. _amd-maxtext-get-started-v26.3:
+.. _amd-maxtext-get-started-v26.4:
 
 Benchmarking
 ============
@@ -169,7 +169,7 @@ benchmark results:
             .. container:: model-doc {{ model.mad_tag }}
 
                The following run commands are tailored to {{ model.model }}.
-               See :ref:`amd-maxtext-model-support-v26.3` to switch to another available model.
+               See :ref:`amd-maxtext-model-support-v26.4` to switch to another available model.
 
                .. rubric:: Download the Docker image and required packages
 
@@ -265,7 +265,7 @@ benchmark results:
 
                        .. code-block:: shell
 
-                          ./primus-cli container --image rocm/jax-training:maxtext-v26.3 \
+                          ./primus-cli container --image {{ docker.pull_tag }} \
                             -- train pretrain \
                             --config examples/maxtext/configs/MI300X/{{ model.primus_config_name }}
 
@@ -301,7 +301,7 @@ benchmark results:
          .. tab-item:: MAD-integrated benchmarking
 
             The following run command is tailored to {{ model.model }}.
-            See :ref:`amd-maxtext-model-support-v26.3` to switch to another available model.
+            See :ref:`amd-maxtext-model-support-v26.4` to switch to another available model.
 
             1. Clone the ROCm Model Automation and Dashboarding (`<https://github.com/ROCm/MAD>`__) repository to a local
                directory and install the required packages on the host machine.
@@ -332,7 +332,7 @@ benchmark results:
          .. tab-item:: Standalone benchmarking
 
             The following commands are optimized for {{ model.model }}. See
-            :ref:`amd-maxtext-model-support-v26.3` to switch to another
+            :ref:`amd-maxtext-model-support-v26.4` to switch to another
             available model. Some instructions and resources might not be
             available for all models and configurations.
 
@@ -452,7 +452,7 @@ benchmark results:
 
             [docker_image] (optional)
                The Docker image to use. If not specified, it defaults to
-               ``rocm/jax-training:maxtext-v26.3``.
+               ``rocm/jax-training:maxtext-v26.4-jax0.9.1-te2.12.0``.
 
             For example, to run a multi-node training benchmark on {{ model.model }}:
 
@@ -477,7 +477,7 @@ benchmark results:
          {% else %}
             .. rubric:: Multi-node training
 
-            For multi-node training examples, choose a model from :ref:`amd-maxtext-model-support-v26.3`
+            For multi-node training examples, choose a model from :ref:`amd-maxtext-model-support-v26.4`
             with an available `multi-node training script <https://github.com/ROCm/MAD/tree/develop/scripts/jax-maxtext/env_scripts>`__.
          {% endif %}
       {% endfor %}
@@ -549,8 +549,8 @@ Example: Profile a model standalone in Docker
    #!/bin/bash
    set -e
 
-   IMAGE="$1"       # Docker image, e.g. rocm/jax-training:maxtext-v26.3
-   TAG="$2"         # Short tag for output folder, e.g. v26.3_llama2_7b
+   IMAGE="$1"       # Docker image, e.g. rocm/jax-training:maxtext-v26.4-jax0.9.1-te2.12.0
+   TAG="$2"         # Short tag for output folder, e.g. v26.4_llama2_7b
    PROFILE_DIR="/path/to/profiles/${TAG}"
 
    mkdir -p "${PROFILE_DIR}"
@@ -646,17 +646,14 @@ Known issues
   Set ``NVTE_CK_IS_V3_ATOMIC_FP32=1`` for production training when using
   real data and input sequence packing (``packing=True``).
 
-- There is a known slight performance regression for DeepSeek-V2-lite
-  (16B) in v26.3. This is being tracked and will be addressed in a future
-  release.
-
-- **JAX 0.9.1 Early Access known issues:**
-
-  - There is a known performance regression for MoE models
-    (DeepSeek-V2-lite and Mixtral-8x7B).
+- There is a known performance regression for Mixtral-8x7B in v26.4.
+  This is being tracked and will be addressed in a future release.
 
-  - The trace viewer in profiling may be missing some information in the
-    flame graph.
+- There is a discrepancy in the loss curve when setting ``packing=False``.
+  It converges at a slightly higher value than previous Docker images. To
+  achieve the same convergence as past Docker images, set
+  ``NVTE_CK_USES_FWD_V3=0`` (using FA v2 for forward instead of FA v3).
+  This is being tracked and will be addressed in a future release.
 
 - Shardy is a new config in JAX 0.6.0. You might get related errors if
   it's not configured correctly. To disable it, set ``shardy=False``

docs/how-to/rocm-for-ai/training/benchmark-docker/previous-versions/jax-maxtext-history.rst

@@ -17,12 +17,28 @@ previous releases of the ``ROCm/jax-training`` Docker image on `Docker Hub <http
      - Components
      - Resources
 
-   * - 26.2 (latest)
+   * - 26.4 (latest)
+     -
+       * ROCm 7.14.0
+       * JAX 0.9.1
+     -
+       * :doc:`Documentation <../jax-maxtext>`
+       * `Docker Hub <https://hub.docker.com/layers/rocm/jax-training/maxtext-v26.4-jax0.9.1-te2.12.0/images/>`__
+
+   * - 26.3
+     -
+       * ROCm 7.2.1
+       * JAX 0.8.2
+     -
+       * :doc:`Documentation <jax-maxtext-v26.3>`
+       * `Docker Hub <https://hub.docker.com/layers/rocm/jax-training/maxtext-v26.3/images/sha256-548381e496fa1367772964341130cc014ee1a9a5ca715ff4ca3aa92960d162c3>`__
+
+   * - 26.2
      -
        * ROCm 7.1.1
        * JAX 0.8.2
      -
-       * :doc:`Documentation <../jax-maxtext>`
+       * :doc:`Documentation <jax-maxtext-v26.2>`
        * `Docker Hub <https://hub.docker.com/layers/rocm/jax-training/maxtext-v26.2/images/sha256-a89643388487b1e2fc6b6ef7bd3c44378c05d217309c977a1c18c72d05ebcaeb>`__
 
    * - 26.1