Executar Jobs de Treinamento do Framework NVIDIA NeMo

O NVIDIA NeMo Framework Launcher é uma ferramenta nativa da nuvem para lançar trabalhos de treinamento completos do NeMo Framework em milhares de GPUs para treinamento de LLM em larga escala. Neste exemplo, usamos o NeMo Framework Launcher para executar o modelo de linguagem gpt3_5b large, a preparação de dados e os estágios de treinamento.

Consulte a documentação da NVIDIA para obter mais detalhes sobre o NeMo e o NeMo Framework Launcher:

Execute uma Carga de Trabalho de Treinamento LLM

Instale o python e execute uma carga de trabalho de treinamento.

  1. Instale o python 3.8 e torne-o o python padrão para o usuário opc.

    Para este Python3 é necessário, mais módulos python listados no requirements.txt. O Oracle Linux 7 (e algumas outras releases do SO) ainda python2.

    $ sudo yum install -y oracle-softwarecollection-release-el7
    $ sudo yum -y install scl-utils rh-python38
    $ scl enable rh-python38 bash
    $ cat <<EOF >> ~/.bashrc
    [ -f /opt/rh/rh-python38/enable ] && source /opt/rh/rh-python38/enable
    EOF
  2. Use os comandos Installation em https://github.com/NVIDIA/NeMo-Megatron-Launcher.
    $ cd /nfs/scratch
    $ git clone https://github.com/NVIDIA/NeMo-Megatron-Launcher.git
    Cloning into 'NeMo-Megatron-Launcher'...
    remote: Enumerating objects: 29018, done.
    remote: Counting objects: 100% (1062/1062), done.
    remote: Compressing objects: 100% (452/452), done.
    remote: Total 29018 (delta 665), reused 898 (delta 564), pack-reused 27956
    Receiving objects: 100% (29018/29018), 27.66 MiB | 14.16 MiB/s, done.
    Resolving deltas: 100% (18124/18124), done.
    
    $ cd NeMo-Megatron-Launcher
    $ pip install -r requirements.txt --user
    $ pip install --upgrade requests --user

    Observação:

    Se usuários diferentes de opc estiverem compartilhando o cluster, você precisará instalar os módulos python para todos os usuários com: sudo pip install -r requirements.txt.

Preparação de Dados

O estágio de preparação de dados executa três tarefas: fazer download do conjunto de dados "a pilha" sem problemas; extrair (descompactar) os dados e pré-processar os dados.

  1. Edite launcher_scripts/conf/config.yaml:
    1. Defina o job stage como data_preparation.
      stages:
        - data_preparation
        #- training
    2. Defina o launcher_scripts_path.
      # Path to NeMo Megatron Launch scripts, should ends with /launcher_scripts
      launcher_scripts_path: /nfs/scratch/NeMo-Megatron-Launcher/launcher_scripts
  2. No diretório launcher_scripts, execute main.py para submeter o job ao Slurm.
    $ cd /nfs/scratch/NeMo-Megatron-Launcher/launcher_scripts
    $ python main.py
    Job nemo-megatron-download_gpt3_pile submission file created at '/nfs/scratch/NeMo-Megatron-Launcher/launcher_scripts/results/download_gpt3_pile/download/nemo-megatron-download_gpt3_pile_submission.sh' 
    . . .
  3. Use o comando squeue do Slurm para o status do job do observador.
    $ squeue
                 JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)
        191_[20-29%30]       gpu nemo-meg      opc PD       0:00      1 (Resources)
            192_[0-29]       gpu nemo-meg      opc PD       0:00      1 (Dependency,Priority)
            193_[0-29]       gpu nemo-meg      opc PD       0:00      1 (Dependency)
                191_19       gpu nemo-meg      opc  R       0:50      1 gpu-permanent-node-517
                191_18       gpu nemo-meg      opc  R       0:56      1 gpu-permanent-node-878
    

    Jobs com status (ST) R indicam que alguns dos jobs estão em execução. Outros jobs mostrados estão aguardando recursos ou aguardando que outros jobs sejam concluídos primeiro (dependências).

    Consulte results/download_gpt3_pile/download para ver a saída de jobs concluídos e em execução. Você também pode encontrar o script bash que foi usado para enviar os trabalhos Slurm. O script e a saída podem ser úteis para solucionar problemas de jobs que não são executados conforme esperado.
Em um cluster de dois nós, as etapas foram executadas:
  • 90 minutos para download
  • 46 minutos para extrair
  • 5 horas 45 minutos para pré-processamento
Os tempos de execução seriam substancialmente menores em um cluster maior em que os fragmentos de dados podem ser paralelizados em até 30 nós.

Treinamento

No estágio de treinamento, você editará os scripts para executar o treinamento de LLM.

  1. Edite launcher_scripts/conf/config.yaml:
    1. Defina o job stage como training.
      stages:
        #- data_preparation
        - training
    2. Adicione variáveis para configurar o NVIDIA NVLink na OCI.
      Na seção env_vars, adicione essas variáveis para configurar o NVIDIA NVLink no OCI. Deixe as variáveis existentes no lugar, mas comente todos os frascos que estamos substituindo por novos valores.
      env_vars:
        TRANSFORMERS_OFFLINE: 0 # (was 1)
        . . .
        RX_QUEUE_LEN: 8192
        IB_RX_QUEUE_LEN: 8192
        UCX_TLS: tcp
        HCOLL_ENABLE_MCAST_ALL: 0
        coll_hcoll_enable: 0
        UCX_NET_DEVICES: ens300
        NCCL_SOCKET_IFNAME: ens300
        NCCL_IB_TIMEOUT: 16
        NCCL_IB_SL: 0
        NCCL_IB_TC: 41
        NCCL_ALGO: Auto  # tree, ring
        NCCL_IB_GID_INDEX: 3
        NCCL_IB_QPS_PER_CONNECTION: 16  # was 4
        NCCL_IB_HCA: \'mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_6,mlx5_7,mlx5_8,mlx5_9,mlx5_10,mlx5_11,mlx5_12,mlx5_14,mlx5_15,mlx5_16,mlx5_17\'
        NCCL_DEBUG: INFO # Logging level for NCCL. Set to "INFO" for debug information
  2. Edite conf/training/gpt3/5b.yaml.
    run:
      time_limit: "6-00:00:00"  # allow the training job to run for 6 days
    
    trainer:
      num_nodes: 2 # (was 16) set to the size of your cluster
    
    model:
      micro_batch_size: 2           # (was 4) change to fit in A100/40GB memory
      tensor_model_parallel_size: 2 # (was 1) change to fit in A100/40GB memory
    
      optim:
        bucket_cap_mb: 200 # (was 400)
  3. No diretório launcher_scripts, execute main.py para submeter o job ao Slurm. Use o comando squeue para verificar se o job está em execução ("R").
    $ cd /nfs/scratch/NeMo-Megatron-Launcher/launcher_scripts
    $ python main.py
    Job nemo-megatron-gpt3_5b submission file created at '/nfs/scratch/NeMo-Megatron-Launcher/launcher_scripts/results/gpt3_5b/nemo-megatron-gpt3_5b_submission.sh'
    Job nemo-megatron-gpt3_5b submitted with Job ID 285
    $ squeue
                 JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)
                   285       gpu nemo-meg      opc  R       0:06      2 gpu-permanent-node-[517,878]
  4. Verifique os arquivos em results/gpt3_5b/* para exibir mensagens de saída e erro e monitorar o andamento dos jobs em execução.
    O treinamento de LLM está funcionando se você vir linhas semelhantes a isso no gpt3_5b_nnn.out:
    Training:   0%|          | 0/75375 [00:00<?]
    Epoch 0: :   0%|          | 0/75375 [00:00<?]
    Epoch 0: :   0%|          | 1/75375 [00:52<1089:12:32]
    Epoch 0: :   0%|          | 1/75375 [00:52<1089:13:02 ... train_step_timing in s=52.00]
    Epoch 0: :   0%|          | 2/75375 [01:33<980:51:18 ... train_step_timing in s=52.00]
    Epoch 0: :   0%|          | 2/75375 [01:33<980:51:33 ... train_step_timing in s=46.80]
    Epoch 0: :   0%|          | 3/75375 [02:15<945:29:05 ... train_step_timing in s=46.80]
    Epoch 0: :   0%|          | 3/75375 [02:15<945:29:14 ... train_step_timing in s=45.20]
    Epoch 0: :   0%|          | 4/75375 [02:57<926:40:09 ... train_step_timing in s=45.20]
    Epoch 0: :   0%|          | 4/75375 [02:57<926:40:16 ... train_step_timing in s=44.30]
    Epoch 0: :   0%|          | 5/75375 [03:38<915:10:14 ... train_step_timing in s=44.30]
    Epoch 0: :   0%|          | 5/75375 [03:38<915:10:20 ... train_step_timing in s=43.70]
    Epoch 0: :   0%|          | 6/75375 [04:20<907:25:47 ... train_step_timing in s=43.70]
    Epoch 0: :   0%|          | 6/75375 [04:20<907:25:52 ... train_step_timing in s=41.60]
    Epoch 0: :   0%|          | 7/75375 [05:01<901:53:34 ... train_step_timing in s=41.60]
    Epoch 0: :   0%|          | 7/75375 [05:01<901:53:38 ... train_step_timing in s=41.60]
    Epoch 0: :   0%|          | 8/75375 [05:43<897:38:17 ... train_step_timing in s=41.60]
    Epoch 0: :   0%|          | 8/75375 [05:43<897:38:21 ... train_step_timing in s=41.50]
    Epoch 0: :   0%|          | 9/75375 [06:24<894:16:56 ... train_step_timing in s=41.50]
    Epoch 0: :   0%|          | 9/75375 [06:24<894:16:59 ... train_step_timing in s=41.50]
    Epoch 0: :   0%|          | 10/75375 [07:05<891:30:50 ... train_step_timing in s=41.50]

    Isso indica que cada etapa de treinamento foi concluída em 41,5 segundos. Tempos de etapa muito mais rápidos podem ser alcançados com mais GPUs no cluster.