.. meta::
   :description: How to train a model using PyTorch for ROCm.
   :keywords: ROCm, AI, LLM, train, PyTorch, torch, Llama, flux, tutorial, docker

****************************************
Training a model with Primus and PyTorch
****************************************

`Primus <https://github.com/AMD-AGI/Primus>`__ is a unified and flexible
LLM training framework designed to streamline training. It streamlines LLM
training on AMD Instinct GPUs using a modular, reproducible configuration paradigm.
Primus now supports the PyTorch torchtitan backend.

.. note::

   For a unified training solution on AMD GPUs with ROCm, the `rocm/pytorch-training
   <https://hub.docker.com/r/rocm/pytorch-training/>`__ Docker Hub registry will be
   deprecated soon in favor of `rocm/primus <https://hub.docker.com/r/rocm/primus>`__.
   The ``rocm/primus`` Docker containers will cover PyTorch training ecosystem frameworks,
   including torchtitan and :doc:`Megatron-LM <primus-megatron>`.

   Primus with the PyTorch torchtitan backend is designed to replace the
   :doc:`ROCm PyTorch training <pytorch-training>` workflow. See
   :doc:`pytorch-training` to see steps to run workloads without Primus.

AMD provides a ready-to-use Docker image for MI355X, MI350X, MI325X, and
MI300X GPUs containing essential components for Primus and PyTorch training
with Primus Turbo optimizations.

.. datatemplate:yaml:: /data/how-to/rocm-for-ai/training/primus-pytorch-benchmark-models.yaml

   {% set dockers = data.dockers %}
   .. tab-set::

   {% for supported_gpus, docker in dockers.items() %}
      .. tab-item:: {{ supported_gpus }}
         :sync: {{ supported_gpus }}

         .. list-table::
            :header-rows: 1

            * - Software component
              - Version

            {% for component_name, component_version in docker.components.items() %}
            * - {{ component_name }}
              - {{ component_version }}
            {% endfor %}
   {% endfor %}

.. _amd-primus-pytorch-model-support-v259:

Supported models
================

The following models are pre-optimized for performance on the AMD Instinct MI325X and MI300X GPUs.
Some instructions, commands, and training recommendations in this documentation might
vary by model -- select one to get started.

.. datatemplate:yaml:: /data/how-to/rocm-for-ai/training/primus-pytorch-benchmark-models.yaml

   {% set model_groups = data.model_groups %}
   .. raw:: html

      <div id="vllm-benchmark-ud-params-picker" class="container-fluid">
         <div class="row gx-0">
            <div class="col-2 me-1 px-2 model-param-head">Model</div>
            <div class="row col-10 pe-0">
      {% for model_group in model_groups %}
               <div class="col-12 px-2 model-param" data-param-k="model-group" data-param-v="{{ model_group.tag }}" tabindex="0">{{ model_group.group }}</div>
      {% endfor %}
            </div>
         </div>

         <div class="row gx-0 pt-1">
            <div class="col-2 me-1 px-2 model-param-head">Variant</div>
            <div class="row col-10 pe-0">
      {% for model_group in model_groups %}
         {% set models = model_group.models %}
         {% for model in models %}
            {% if models|length % 3 == 0 %}
               <div class="col-4 px-2 model-param" data-param-k="model" data-param-v="{{ model.mad_tag }}" data-param-group="{{ model_group.tag }}" tabindex="0">{{ model.model }}</div>
            {% else %}
               <div class="col-6 px-2 model-param" data-param-k="model" data-param-v="{{ model.mad_tag }}" data-param-group="{{ model_group.tag }}" tabindex="0">{{ model.model }}</div>
            {% endif %}
         {% endfor %}
      {% endfor %}
            </div>
         </div>
      </div>

.. seealso::

   For additional workloads, including Llama 3.3, Llama 3.2, Llama 2, GPT OSS, Qwen, and Flux models,
   see the documentation :doc:`pytorch-training` (without Primus)

.. _amd-primus-pytorch-performance-measurements-v259:

System validation
=================

Before running AI workloads, it's important to validate that your AMD hardware is configured
correctly and performing optimally.

If you have already validated your system settings, including aspects like NUMA auto-balancing, you
can skip this step. Otherwise, complete the procedures in the :ref:`System validation and
optimization <rocm-for-ai-system-optimization>` guide to properly configure your system settings
before starting training.

To test for optimal performance, consult the recommended :ref:`System health benchmarks
<rocm-for-ai-system-health-bench>`. This suite of tests will help you verify and fine-tune your
system's configuration.

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 dockers = data.dockers %}

   Use the following command to pull the Docker image from Docker Hub.

   .. tab-set::

      {% for supported_gpus, docker in dockers.items() %}
      .. tab-item:: {{ supported_gpus }}
         :sync: {{ supported_gpus }}

         .. code-block:: shell

            docker pull {{ docker.pull_tag }}
      {% endfor %}

Run training
============

Once the setup is complete, choose between the following two workflows to start benchmarking training.
For fine-tuning workloads and multi-node training examples, see :doc:`pytorch-training` (without Primus).
For best performance on MI325X, MI350X, and MI355X GPUs, you might need to
tweak some configurations (such as batch sizes).

.. datatemplate:yaml:: /data/how-to/rocm-for-ai/training/primus-pytorch-benchmark-models.yaml

   {% set dockers = data.dockers %}
   {% set model_groups = data.model_groups %}

   .. tab-set::

      .. tab-item:: MAD-integrated benchmarking

   {% for model_group in model_groups %}
      {% for model in model_group.models %}

         .. container:: model-doc {{ model.mad_tag }}

            The following run command is tailored to {{ model.model }}.
            See :ref:`amd-primus-pytorch-model-support-v259` 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.

               .. code-block:: shell

                  git clone https://github.com/ROCm/MAD
                  cd MAD
                  pip install -r requirements.txt

            2. For example, use this command to run the performance benchmark test on the {{ model.model }} model
               using one node with the {{ model.precision }} data type on the host machine.

               .. code-block:: shell

                  export MAD_SECRETS_HFTOKEN="your personal Hugging Face token to access gated models"
                  madengine run \
                      --tags {{ model.mad_tag }} \
                      --keep-model-dir \
                      --live-output \
                      --timeout 28800

               MAD launches a Docker container with the name
               ``container_ci-{{ model.mad_tag }}``. The latency and throughput reports of the
               model are collected in ``~/MAD/perf.csv``.

               .. note::

                  Currently, Primus torchtitan models are run with Primus Turbo
                  enabled for enhanced performance. To disable Primus Turbo,
                  modify respective configuration file
                  ``scripts/primus/pytorch_train/primus_torchtitan_scripts/llama3_[8B|70B]-[BF16|FP8].yaml``.

      {% endfor %}
   {% endfor %}

      .. tab-item:: Primus benchmarking

   {% for model_group in model_groups %}
      {% for model in model_group.models %}

         .. container:: model-doc {{ model.mad_tag }}

            The following run commands are tailored to {{ model.model }}.
            See :ref:`amd-primus-pytorch-model-support-v259` to switch to another available model.

            .. rubric:: Download the Docker image and required packages

            1. Pull the appropriate Docker image for your AMD GPU architecture from Docker Hub.

               .. tab-set::

                  {% for supported_gpus, docker in dockers.items() %}
                  .. tab-item:: {{ supported_gpus }}
                     :sync: {{ supported_gpus }}

                     .. code-block:: shell

                        docker pull {{ docker.pull_tag }}
                  {% endfor %}

            2. Run the Docker container.

               .. tab-set::

                  {% for supported_gpus, docker in dockers.items() %}
                  .. tab-item:: {{ supported_gpus }}
                     :sync: {{ supported_gpus }}

                     .. code-block:: shell

                        docker run -it \
                            --device /dev/dri \
                            --device /dev/kfd \
                            --network host \
                            --ipc host \
                            --group-add video \
                            --cap-add SYS_PTRACE \
                            --security-opt seccomp=unconfined \
                            --privileged \
                            -v $HOME:$HOME \
                            -v $HOME/.ssh:/root/.ssh \
                            --shm-size 64G \
                            --name training_env \
                            {{ docker.pull_tag }}
                  {% endfor %}

               Use these commands if you exit the ``training_env`` container and need to return to it.

               .. code-block:: shell

                  docker start training_env
                  docker exec -it training_env bash

            .. rubric:: Prepare training datasets and dependencies

            The following benchmarking examples require downloading models and datasets
            from Hugging Face. To ensure successful access to gated repos, set your
            ``HF_TOKEN``.

            .. code-block:: shell

               export HF_TOKEN=$your_personal_hugging_face_access_token

            .. rubric:: Pretraining

            To get started, navigate to the ``Primus`` directory in your container.

            .. code-block::

               cd /workspace/Primus

            Now, to start the pretraining benchmark, use the ``run_pretrain.sh`` script
            included with Primus with the appropriate options.

            .. rubric:: Benchmarking examples

            .. container:: model-doc primus_pyt_train_llama-3.1-8b

               Use the following command to run train Llama 3.1 8B with BF16 precision using Primus torchtitan.

               .. tab-set::

                  .. tab-item:: MI355X and MI350X
                     :sync: MI355X and MI300X

                     .. code-block:: shell

                        EXP=examples/torchtitan/configs/llama3.1_8B-BF16-pretrain.yaml \
                        bash examples/run_pretrain.sh \
                            --metrics.enable_tensorboard false \
                            --profiling.enable_profiling false \
                            --training.batch_size 5

                  .. tab-item:: MI325X
                     :sync: MI325X

                     .. code-block:: shell

                        EXP=examples/torchtitan/configs/llama3.1_8B-BF16-pretrain.yaml \
                        bash examples/run_pretrain.sh \
                            --metrics.enable_tensorboard false \
                            --profiling.enable_profiling false \
                            --training.batch_size 6

                  .. tab-item:: MI300X
                     :sync: MI325X and MI300X

                     .. code-block:: shell

                        EXP=examples/torchtitan/configs/llama3.1_8B-BF16-pretrain.yaml \
                        bash examples/run_pretrain.sh \
                            --metrics.enable_tensorboard false \
                            --profiling.enable_profiling false \
                            --training.batch_size 4


               To train Llama 3.1 8B with FP8 precision, use the following command.

               .. tab-set::

                  .. tab-item:: MI355X and MI350X
                     :sync: MI355X and MI300X

                     .. code-block:: shell

                        EXP=examples/torchtitan/configs/llama3.1_8B-BF16-pretrain.yaml \
                        bash examples/run_pretrain.sh \
                            --metrics.enable_tensorboard false \
                            --profiling.enable_profiling false \
                            --training.batch_size 8

                  .. tab-item:: MI325X
                     :sync: MI325X

                     .. code-block:: shell

                        EXP=examples/torchtitan/configs/llama3.1_8B-FP8-pretrain.yaml \
                        bash examples/run_pretrain.sh \
                            --metrics.enable_tensorboard false \
                            --profiling.enable_profiling false \
                            --training.batch_size 7

                  .. tab-item:: MI300X
                     :sync: MI325X and MI300X

                     .. code-block:: shell

                        EXP=examples/torchtitan/configs/llama3.1_8B-FP8-pretrain.yaml \
                        bash examples/run_pretrain.sh \
                            --metrics.enable_tensorboard false \
                            --profiling.enable_profiling false \
                            --training.batch_size 5

            .. container:: model-doc primus_pyt_train_llama-3.1-70b

               Use the following command to run train Llama 3.1 70B with BF16 precision using Primus torchtitan.

               .. tab-set::

                  .. tab-item:: MI355X and MI350X
                     :sync: MI355X and MI300X

                     .. code-block:: shell

                        EXP=examples/torchtitan/configs/llama3.1_70B-BF16-pretrain.yaml \
                        bash examples/run_pretrain.sh \
                            --metrics.enable_tensorboard false \
                            --profiling.enable_profiling false \
                            --training.batch_size 8

                  .. tab-item:: MI325X
                     :sync: MI325X

                     .. code-block:: shell

                        EXP=examples/torchtitan/configs/llama3.1_70B-BF16-pretrain.yaml \
                        bash examples/run_pretrain.sh \
                            --metrics.enable_tensorboard false \
                            --profiling.enable_profiling false \
                            --training.batch_size 6

                  .. tab-item:: MI300X
                     :sync: MI325X and MI300X

                     .. code-block:: shell

                        EXP=examples/torchtitan/configs/llama3.1_70B-BF16-pretrain.yaml \
                        bash examples/run_pretrain.sh \
                            --metrics.enable_tensorboard false \
                            --profiling.enable_profiling false \
                            --training.batch_size 4

               To train Llama 3.1 70B with FP8 precision, use the following command.

               .. tab-set::

                  .. tab-item:: MI355X and MI350X
                     :sync: MI355X and MI300X

                     .. code-block:: shell

                        EXP=examples/torchtitan/configs/llama3.1_70B-FP8-pretrain.yaml \
                        bash examples/run_pretrain.sh \
                            --metrics.enable_tensorboard false \
                            --profiling.enable_profiling false \
                            --training.batch_size 6

                  .. tab-item:: MI325X
                     :sync: MI325X

                     .. code-block:: shell

                        EXP=examples/torchtitan/configs/llama3.1_70B-FP8-pretrain.yaml \
                        bash examples/run_pretrain.sh \
                            --metrics.enable_tensorboard false \
                            --profiling.enable_profiling false \
                            --training.batch_size 5

                  .. tab-item:: MI300X
                     :sync: MI325X and MI300X

                     .. code-block:: shell

                        EXP=examples/torchtitan/configs/llama3.1_70B-FP8-pretrain.yaml \
                        bash examples/run_pretrain.sh \
                            --metrics.enable_tensorboard false \
                            --profiling.enable_profiling false \
                            --training.batch_size 3
      {% endfor %}
   {% endfor %}

      .. tab-item:: Standalone torchtitan benchmarking

   {% for model_group in model_groups %}
      {% for model in model_group.models %}

         .. container:: model-doc {{ model.mad_tag }}

            The following run commands are tailored to {{ model.model }}.
            See :ref:`amd-primus-pytorch-model-support-v259` to switch to another available model.

            .. rubric:: Download the Docker image and required packages

            1. Pull the appropriate Docker image for your AMD GPU architecture from Docker Hub.

               .. tab-set::

                  {% for supported_gpus, docker in dockers.items() %}
                  .. tab-item:: {{ supported_gpus }}
                     :sync: {{ supported_gpus }}

                     .. code-block:: shell

                        docker pull {{ docker.pull_tag }}
                  {% endfor %}

            2. Run the Docker container.

               .. tab-set::

                  {% for supported_gpus, docker in dockers.items() %}
                  .. tab-item:: {{ supported_gpus }}
                     :sync: {{ supported_gpus }}

                     .. code-block:: shell

                        docker run -it \
                            --device /dev/dri \
                            --device /dev/kfd \
                            --network host \
                            --ipc host \
                            --group-add video \
                            --cap-add SYS_PTRACE \
                            --security-opt seccomp=unconfined \
                            --privileged \
                            -v $HOME:$HOME \
                            -v $HOME/.ssh:/root/.ssh \
                            --shm-size 64G \
                            --name training_env \
                            {{ docker.pull_tag }}
                  {% endfor %}

               Use these commands if you exit the ``training_env`` container and need to return to it.

               .. code-block:: shell

                  docker start training_env
                  docker exec -it training_env bash

            3. Navigate to the ``torchtitan`` workspace directory.

               .. code-block:: shell

                  cd /workspace/torchtitan

            .. rubric:: Download the tokenizer

            1. The following benchmarking examples require downloading models and datasets
               from Hugging Face. To ensure successful access to gated repos, set your
               ``HF_TOKEN``.

               .. code-block:: shell

                  export HF_TOKEN=$your_personal_hugging_face_access_token

            2. Download the tokenizer for your model.

               .. container:: model-doc {{ model.mad_tag }}

                  .. code-block:: shell

                     python3 scripts/download_tokenizer.py \
                        --repo_id {{ model.model_repo }} \
                        --tokenizer_path "original" \
                        --hf_token=${HF_TOKEN}

            .. rubric:: Pretraining examples

            Run the training script with the appropriate configuration file.

            For train with BF16 precicion, use the following command:

            .. container:: model-doc {{ model.mad_tag }}

               .. code-block:: shell

                  CONFIG_FILE={{ model.config_file.bf16 }} \
                  .run_train.sh

            For train with BF16 precicion, use the following command:

            .. container:: model-doc {{ model.mad_tag }}

               .. code-block:: shell

                  CONFIG_FILE={{ model.config_file.fp8 }} \
                  .run_train.sh
      {% endfor %}
   {% endfor %}

Known issues
============

PyTorch Profiler may produce inaccurate traces when CPU activity profiling is enabled.


Further reading
===============

- For an introduction to Primus, see `Primus: A Lightweight, Unified Training
  Framework for Large Models on AMD GPUs <https://rocm.blogs.amd.com/software-tools-optimization/primus/README.html>`__.

- To learn more about MAD and the ``madengine`` CLI, see the `MAD usage guide <https://github.com/ROCm/MAD?tab=readme-ov-file#usage-guide>`__.

- To learn more about system settings and management practices to configure your system for
  AMD Instinct MI300X Series GPUs, see `AMD Instinct MI300X system optimization <https://instinct.docs.amd.com/projects/amdgpu-docs/en/latest/system-optimization/mi300x.html>`_.

- For a list of other ready-made Docker images for AI with ROCm, see
  `AMD Infinity Hub <https://www.amd.com/en/developer/resources/infinity-hub.html#f-amd_hub_category=AI%20%26%20ML%20Models>`_.

Previous versions
=================

See :doc:`previous-versions/pytorch-training-history` to find documentation for previous releases
of the ``ROCm/pytorch-training`` Docker image.
