diff --git a/docs/how-to/rocm-for-ai/inference/benchmark-docker/vllm.rst b/docs/how-to/rocm-for-ai/inference/benchmark-docker/vllm.rst
index 70121c67f..38a5f8200 100644
--- a/docs/how-to/rocm-for-ai/inference/benchmark-docker/vllm.rst
+++ b/docs/how-to/rocm-for-ai/inference/benchmark-docker/vllm.rst
@@ -138,13 +138,12 @@ To test for optimal performance, consult the recommended :ref:`System health ben
 <rocm-for-ai-system-health-bench>`. This suite of tests will help you verify and fine-tune your
 system's configuration.
 
+Pull the Docker image
+=====================
+
 .. datatemplate:yaml:: /data/how-to/rocm-for-ai/inference/vllm-benchmark-models.yaml
 
    {% set docker = data.dockers[0] %}
-   {% set model_groups = data.model_groups %}
-
-   Pull the Docker image
-   =====================
 
    Download the `ROCm vLLM Docker image <{{ docker.docker_hub_url }}>`_.
    Use the following command to pull the Docker image from Docker Hub.
@@ -153,8 +152,13 @@ system's configuration.
 
       docker pull {{ docker.pull_tag }}
 
-   Benchmarking
-   ============
+Benchmarking
+============
+
+.. datatemplate:yaml:: /data/how-to/rocm-for-ai/inference/vllm-benchmark-models.yaml
+
+   {% set docker = data.dockers[0] %}
+   {% set model_groups = data.model_groups %}
 
    Once the setup is complete, choose between two options to reproduce the
    benchmark results:
diff --git a/docs/how-to/rocm-for-ai/training/benchmark-docker/jax-maxtext.rst b/docs/how-to/rocm-for-ai/training/benchmark-docker/jax-maxtext.rst
index eb56f4dce..eec785b7b 100644
--- a/docs/how-to/rocm-for-ai/training/benchmark-docker/jax-maxtext.rst
+++ b/docs/how-to/rocm-for-ai/training/benchmark-docker/jax-maxtext.rst
@@ -25,7 +25,7 @@ It includes the following software components:
       {% for docker in dockers %}
       {% set jax_version = docker.components["JAX"] %}
 
-      .. tab-item:: JAX {{ jax_version }}
+      .. tab-item:: ``{{ docker.pull_tag }}``
          :sync: {{ docker.pull_tag }}
 
          .. list-table::
@@ -132,6 +132,28 @@ This Docker image is optimized for specific model configurations outlined
 as follows. Performance can vary for other training workloads, as AMD
 doesn’t validate configurations and run conditions outside those described.
 
+Pull the Docker image
+---------------------
+
+Use the following command to pull the Docker image from Docker Hub.
+
+.. datatemplate:yaml:: /data/how-to/rocm-for-ai/training/jax-maxtext-benchmark-models.yaml
+
+   {% set dockers = data.dockers %}
+   .. tab-set::
+
+      {% for docker in dockers %}
+      {% set jax_version = docker.components["JAX"] %}
+
+      .. tab-item:: JAX {{ jax_version }}
+         :sync: {{ docker.pull_tag }}
+
+         .. code-block:: shell
+
+            docker pull {{ docker.pull_tag }}
+
+      {% endfor %}
+
 .. _amd-maxtext-multi-node-setup-v257:
 
 Multi-node configuration
diff --git a/docs/how-to/rocm-for-ai/training/benchmark-docker/primus-megatron.rst b/docs/how-to/rocm-for-ai/training/benchmark-docker/primus-megatron.rst
index 853d24395..65ac5e50c 100644
--- a/docs/how-to/rocm-for-ai/training/benchmark-docker/primus-megatron.rst
+++ b/docs/how-to/rocm-for-ai/training/benchmark-docker/primus-megatron.rst
@@ -105,21 +105,26 @@ system's configuration.
 
 .. _mi300x-amd-primus-megatron-lm-training:
 
+Environment setup
+=================
+
 .. datatemplate:yaml:: /data/how-to/rocm-for-ai/training/primus-megatron-benchmark-models.yaml
 
    {% set dockers = data.dockers %}
       {% set docker = dockers[0] %}
 
-   Environment setup
-   =================
-
    Use the following instructions to set up the environment, configure the script to train models, and
    reproduce the benchmark results on MI300X series GPUs with the ``{{ docker.pull_tag }}`` image.
 
    .. _amd-primus-megatron-lm-requirements:
 
-   Download the Docker image
-   -------------------------
+Pull the Docker image
+=====================
+
+.. datatemplate:yaml:: /data/how-to/rocm-for-ai/training/primus-megatron-benchmark-models.yaml
+
+   {% set dockers = data.dockers %}
+      {% set docker = dockers[0] %}
 
    1. Use the following command to pull the Docker image from Docker Hub.
 
diff --git a/docs/how-to/rocm-for-ai/training/benchmark-docker/primus-pytorch.rst b/docs/how-to/rocm-for-ai/training/benchmark-docker/primus-pytorch.rst
index b2bd5fb87..5c99776ee 100644
--- a/docs/how-to/rocm-for-ai/training/benchmark-docker/primus-pytorch.rst
+++ b/docs/how-to/rocm-for-ai/training/benchmark-docker/primus-pytorch.rst
@@ -104,22 +104,25 @@ This Docker image is optimized for specific model configurations outlined
 below. Performance can vary for other training workloads, as AMD
 doesn’t test configurations and run conditions outside those described.
 
+Pull the Docker image
+=====================
+
 .. datatemplate:yaml:: /data/how-to/rocm-for-ai/training/primus-pytorch-benchmark-models.yaml
 
    {% set unified_docker = data.dockers[0] %}
 
-   Pull the Docker image
-   =====================
-
    Use the following command to pull the `Docker image <{{ unified_docker.docker_hub_url }}>`_ from Docker Hub.
 
    .. code-block:: shell
 
       docker pull {{ unified_docker.pull_tag }}
 
-   Run training
-   ============
+Run training
+============
 
+.. datatemplate:yaml:: /data/how-to/rocm-for-ai/training/primus-pytorch-benchmark-models.yaml
+
+   {% set unified_docker = data.dockers[0] %}
    {% set model_groups = data.model_groups %}
 
    Once the setup is complete, choose between the following two workflows to start benchmarking training.