第 4 章 NxD対応済モデルの学習編
本章では以下を仮定します。
- 第 3 章での計算機クラスタ構築を完了している
問題 (19-32)
いよいよ計算機クラスタのヘッドノードに接続して、環境構築、およびモデルの学習を行っていきます。
-
ssh でヘッドノードにログインしてください。
- ヘッドノードのIPアドレスは、前章の CloudShell 上で
aws cloudformation describe-stacks --stack-name $PCLUSTER_NAME --query "Stacks[0].Outputs"を実行するか、あるいはAWSマネジメントコンソールのEC2のインスタンス一覧から確認できます。
- ヘッドノードのIPアドレスは、前章の CloudShell 上で
-
公式のNxDTインストール手順に従って、Python仮想環境を
~/env/に作成した上で、その中にneuronx-distributed-training(NxDT) をインストールしてください。- NxDTとは、一言で言うと「NeuronチップでLLMの学習するならとりあえずこれを入れましょう！コンフィグファイルだけ書いたら学習できて便利ですよ！」という感じのライブラリです。公式説明
-
今回の学習の起点となるベースモデルを、HF Models からダウンロードしてください。
-
環境変数を設定
export HF_MODEL_NAME=weblab-GENIAC/Tanuki-8B-dpo-v1.0 export ORIGINAL_MODEL_PATH=/fsx/models/Tanuki-8B-dpo-v1.0 -
ダウンロード
huggingface-cli download $HF_MODEL_NAME --local-dir $ORIGINAL_MODEL_PATHモデルの重み、トークナイザー、各種configファイルがダウンロードされます。
-
-
slurm の基本的な使い方（習熟済の方はスキップ可）
slurm は、計算機クラスタ向けジョブスケジューラーです。基本的な使い方は
sbatch hoge.shです。このコマンドにより、コンピュートノードが自動で立ち上がり、コンピュートノード上で hoge.sh が実行されます。squeueでジョブの状態を確認でき、sinfoでノードの状態を確認できます。-
sbatch --wrap="echo Hello && pwd && which python && sleep 300 && echo Done"を実行します。直ちにSubmitted batch job xxxと表示されればジョブ投入成功です。xxxの部分は「ジョブ名」です。 -
ジョブが完了するまでの間に
squeueコマンドを実行すると、以下のような結果が返ってきます。JOBID PARTITION NAME USER ST TIME NODES NODELIST(REASON) 1 compute1 wrap ubuntu CF 1:16 1 compute1-dy-queue1-i1- ST と記載された列がジョブの状態 (status) です。主に以下の状態があります:
-
PD: 待機中 -
CF: ジョブ実行準備中 (configuring) -
R: ジョブ実行中 (running)
-
- ST と記載された列がジョブの状態 (status) です。主に以下の状態があります:
-
ジョブが完了するまでの間に
sinfoを実行すると、以下のような結果が返ってきます。PARTITION AVAIL TIMELIMIT NODES STATE NODELIST compute1* up infinite 1 mix# compute1-dy-queue1-gpu-1 compute1* up infinite 27 idle~ compute1-dy-queue1-gpu-2,compute1-dy-queue1-i1-[1-24],compute1-dy-queue1-mem-[1-2]- STATE が
upとなっているのが起動中のノードです。idleとなっているのは起動していません。
- STATE が
-
ssh <sinfoで表示されたノード名>とすると、稼働中のコンピュートノードに ssh 接続できます（上の例の場合、例えばcompute1-dy-queue1-gpu-1やcompute1-dy-queue1-i1-6がノード名です）。 -
稼働中のコンピュートノードは、AWSマネジメントコンソール上でEC2のインスタンス一覧画面にも表示されます。
-
ジョブを途中で取り止めたい時は
scancel <ジョブ名>を実行してください。（上記で実行したジョブは300秒で勝手に終了しますが） -
ジョブ実行中の出力は
slurm-<ジョブ名>.out（例:slurm-1.out）というファイルに書き出されています。このファイルはジョブ実行開始時に（状態がRになるタイミングで）作成され、リアルタイムに更新されます。今回の場合、ジョブ完了時点での中身は以下のようになるはずです。Hello /home/ubuntu env /home/ubuntu/env/bin/python Done -
pwdの結果が/home/ubuntuで、lsの結果にenvと表示されていることからわかるように、/home/ubuntuはマウントされています。また/fsxもマウントされています。さらに、which pythonの結果が/home/ubuntu/env/bin/pythonとなっていることからわかるように、環境変数も丸ごと引き継がれており、結果としてコンピュートノード側でも Python 仮想環境が有効化されています。
💡 slurm の FAQ
- 投入したジョブがいくら待っても実行中 (
R) にならない- ジョブの STATE が
PD(PENDING) になっている場合、scontrol show job <JOBID>のReason=行に保留理由が詳しく表示されます。Resources、Dependency、QOSMaxWallDurationLimitなど、出力された理由に応じて対処してください。
- ジョブの STATE が
- ジョブを投入してから
squeueしたけどジョブが表示されない- 投入したジョブが直ちに実行され、非常に短時間で終了した可能性があります。
slurm-<JOBID>.outなどのログが生成されていれば、その可能性が高いです。sacct -j <JOBID>でステータス (COMPLETED,FAILEDなど) を確認できます。
- 投入したジョブが直ちに実行され、非常に短時間で終了した可能性があります。
-
sbatchの主な実行オプションを知りたい-
主なものは以下です
-
-J, --job-name=<name>: ジョブ名 -
-w, --nodelist=<host1[,host2,...]>: 実行ノードを直接指定 -
--dependency=afterok:<id>: ジョブ依存関係 -
-o, --output=<file>/-e, --error=<file>: 標準出力 / 標準エラーのログファイル名（デフォルトはslurm-%j.out） -
--mail-type=<BEGIN|END|FAIL|ALL>&--mail-user=<email>: メール通知設定 -
--export=<ALL|NONE|VAR1,VAR2=val,...>: 環境変数のエクスポート -
--exclusive: ノードを専有して実行 (同ノード上で他ジョブと共有しない)
-
-
これらのオプションをコマンドラインで指定する代わりに、実行するシェルスクリプトの中に書き込む形で指定することもできます。以下のように
#SBATCHで始まる行を書くことで、sbatch 側が認識してくれます:#!/bin/bash #SBATCH --job-name=test #SBATCH --partition=cpu #SBATCH --nodes=1 #SBATCH --time=00:10:00 #SBATCH --output=slurm-%j.out echo "Hello from node $SLURMD_NODENAME"
-
-
Trainium 上でモデルパラレル学習を行う上で都合の良い形に、チェックポイントの形式を変換する必要があります。
-
学習の起点となるチェックポイントを HuggingFace Transformers (HF) 形式から NeuronX Distributed (NxD) 形式に変換します。
-
NxDTのリポジトリにあるチェックポイント変換スクリプトを使用するため、NxDTのリポジトリをcloneします。
cd ~ git clone https://github.com/aws-neuron/neuronx-distributed-training.git # 変換スクリプトの場所: neuronx-distributed-training/examples/checkpoint_converter_scripts/checkpoint_converter.py -
環境変数を設定します。
export MODEL_CONFIG_PATH=$ORIGINAL_MODEL_PATH/config.json export CHECKPOINT_DIR=/fsx/checkpoints/tanuki-8b/pretrained_model/ export TP_DEGREE=32 export PP_DEGREE=1 export KV_REPLICATOR=4TP_DEGREEはテンソル並列数、PP_DEGREEはパイプライン並列数です（「テンソル並列」「パイプライン並列」について詳しくは後の章をご参照ください）。KV_REPLICATORは「KVヘッド重みの複製回数」のことで、Qヘッド数よりもKVヘッド数が小さなモデル（=Grouped Query Attention (GQA) が使用されているモデル）において、両者を揃えてテンソル分割可能にするために設定するものです。今回のモデルでは Qヘッド数は 32、KV ヘッド数が 8 なので（config.json 中で"num_attention_heads": 32, "num_key_value_heads": 8と確認できます）、32/8 = 4 を設定しています（「KV_REPLICATOR」について詳しくは後の章をご参照ください）。 -
変換スクリプト中で使用される NxD ライブラリ中のコードを、非対話環境向けに一部修正します。
-
pip show neuronx-distributedにより、NxD ライブラリのソースが格納されているディレクトリを確認してください。 -
（上記ディレクトリ）/scripts/checkpoint_converter.pyの中を以下のように編集してください。-
変更前
token = getpass("Enter your Hugging Face API token: (If you don't have one, create it at https://huggingface.co/settings/tokens): ") login(token=token) -
変更後
token = None
-
-
-
以下
convert_checkpoint_from_hf_to_nxd.shを保存した上で、sbatch convert_checkpoint_from_hf_to_nxd.shを実行してください（変換スクリプトの仕様上、変換はヘッドノードではなくコンピュートノード上で行う必要があります）。convert_checkpoint_from_hf_to_nxd.sh
#!/bin/bash echo "HF_MODEL_NAME: $HF_MODEL_NAME" echo "MODEL_CONFIG_PATH: $MODEL_CONFIG_PATH" echo "CHECKPOINT_DIR: $CHECKPOINT_DIR" echo "TP_DEGREE: $TP_DEGREE" echo "PP_DEGREE: $PP_DEGREE" echo "KV_REPLICATOR: $KV_REPLICATOR" cd ~/neuronx-distributed-training/examples/checkpoint_converter_scripts python3 checkpoint_converter.py \ --model_style hf \ --hw_backend trn1 \ --hf_model_name $HF_MODEL_NAME \ --output_dir $CHECKPOINT_DIR \ --save_xser True \ --config $MODEL_CONFIG_PATH \ --tp_size $TP_DEGREE \ --pp_size $PP_DEGREE \ --n_layers 32 \ --kv_size_multiplier $KV_REPLICATOR \ --qkv_linear True \ --fuse_qkv True \ --convert_from_full_state-
checkpoint_converter.pyのコマンドラインオプションについて-
--qkv_linear: 変換対象がGQAモデルである（i.e. Qヘッド数とKVヘッド数が等しくない）場合に True としてください（この場合、内部でGQAQKVLinearモジュールが使用されます。これがオプション名の由来です）。 -
--fuse_qkv: 上記GQAQKVLinearモジュールを使用する場合に、Q, K, Vの計算に用いる3つのパラメータテンソルを、結合された単一のテンソルとして保持するか否かのオプションです（結合しておいた方が計算効率は高いです。デフォルトはTrueです）。
-
成功すると、output_dir の中に、モデル並列数の数（今回は32個）に分割された重みファイルが格納されます。
-
補足（余談）
- チェックポイント変換スクリプトには「ローカルの重みから変換するモード」と「HFからの直接変換モード」があります。上記手順では後者のモードを使用しています。
- ただし後者のモードは、任意の状況で HF TOKEN を対話的に要求してくる仕様となっています。今回は「sbatch により起動されるコンピュートノード」という非対話環境で動作させたいため、上記の修正が必要となっています。
- なお、前者のモードで動かす場合は、ローカルの重みが「単一のファイル」に纏まっている必要があります。既に DL 済みの
weblab-GENIAC/Tanuki-8B-dpo-v1.0の重みは、複数の .safetensors ファイルに分かれて保存されているため、これをまず「単一のファイル」に纏めるための変換処理が必要となります。面倒……
-
-
想定通りにチェックポイント変換が実施されたかどうかを確認するため、以下を確認してください：
- 分割された後の重みファイル
dp_rank_*_tp_rank_*_pp_rank_*.ptが、(DP×TP×PP=)32個存在することを確かめてください。 - 重みファイルの格納されたディレクトリの容量がおよそ21GBであることを、
duコマンドを用いて確認してください。- なお、HF形式の時点での重さ合計はおよそ14GBです。分割前と比べて分割後の方が合計サイズが重くなっていますが、主な原因は、
KV_REPLICATORの設定によるKVヘッドの重み複製です。
- なお、HF形式の時点での重さ合計はおよそ14GBです。分割前と比べて分割後の方が合計サイズが重くなっていますが、主な原因は、
- 分割された後の重みファイル
-
以下のコマンドを実行し、教師データセットを準備してください。
-
環境変数設定
export DATA_DIR=/fsx/examples_datasets/llama3_8b -
コマンド
mkdir -p ${DATA_DIR} && cd ${DATA_DIR} aws s3 cp s3://neuron-s3/training_datasets/llama/sft/training.jsonl . --no-sign-request aws s3 cp s3://neuron-s3/training_datasets/llama/sft/validation.jsonl . --no-sign-request
ここでDLされるデータセットは、各行が
{"input": string, "output": string, "category": string}の形式になっています。training.jsonlから数行だけ抜粋したものを以下に記します。{"input": "Identify which car manufacturer is British or American: Jaguar, GMC", "output": "GMC is American, Jaguar is British", "category": "classification"} {"input": "Best place to visit in West bengal", "output": "1. Darjeeling \n2, Kolkata", "category": "brainstorming"} {"input": "What are the benefits of using genetically modified crops in agriculture?", "output": "There are many benefits to using genetically modified crops in agriculture. The benefits include making the crops hardier with greater resistance to diseases and droughts. The resulting crops can also make food more nutritious and taste better. An additional environmental benefit of genetically modified crops is that they have the potential to reduce resources needed to grow, including water and fertilizer.", "category": "creative_writing"} -
-
slurm の並列実行オプションの理解（習熟済の方はスキップ可）
-
以降では、slurm を使って同一プログラムを複数のノードで（あるいは単一ノード内でも複数タスクを）並列的に走らせる必要があります。
-
ジョブ・ノード・タスク・スレッドは階層関係にあります。ジョブとノードは1:Nの関係にあり、ノードとタスクも1:Nの関係、タスクとスレッドも1:Nの関係にあります。
-
以下のオプションがあります
オプション 意味 --nodes Nジョブ全体のノード数を指定 --ntasks Nジョブ全体のタスク数を指定 --tasks-per-node N1ノードあたりのタスク数を指定 --cpus-per-task N1タスクあたりのスレッド数を指定
-
-
AOTコンパイルの開始
-
conf ディレクトリ内に yaml ファイルが多数あることを確認してください。この中に、今回の学習の設定内容が記載された、以下の内容の
hf_tanuki_8B_SFT_config.yamlを作成してください。hf_tanuki_8B_SFT_config.yaml
name: hf_tanuki model_source: hf seed: 1234 trainer: devices: 32 num_nodes: 1 max_epochs: -1 # PTL default. In practice, max_steps will be reached first. max_steps: 100 # consumed_samples = global_step * micro_batch_size * data_parallel_size * accumulate_grad_batches log_every_n_steps: 10 val_check_interval: 500 # we do not want val to run during training, hence setting it at a high number check_val_every_n_epoch: null num_sanity_val_steps: 0 limit_val_batches: 0.0 limit_test_batches: 0.0 gradient_clip_val: 1.0 exp_manager: log_local_rank_0_only: True # reduce file system access pressure create_tensorboard_logger: True explicit_log_dir: null exp_dir: /fsx/nemo_experiments/ # 実験結果（含チェックポイント）が保存されるディレクトリ name: hf_tanuki resume_if_exists: True resume_ignore_no_checkpoint: True create_checkpoint_callback: True checkpoint_callback_params: monitor: step save_top_k: 1 mode: max save_last: False filename: 'hf_tanuki_8B_SFT--{step}-{consumed_samples}' model_parallel_size: ${multiply:${distributed_strategy.tensor_model_parallel_size}, ${distributed_strategy.pipeline_model_parallel_size}} every_n_train_steps: 50 # このステップ数ごとにチェックポイントを保存 use_master_weights_in_ckpt: False log_parameter_norm: True # Logs parameter norm across model parallel ranks log_gradient_norm: True # Logs gradient norm across model parallel ranks enable_recovery_time_instrumentation: False # default to not printing the detailing timing for recovery save_xser: True load_xser: True save_bf16: True async_checkpointing: False # default to not use async checkpointing resume_from_checkpoint: ${oc.env:CHECKPOINT_DIR} # manually set the checkpoint file to load from. [ModelAlignment change required] ckpt_ptl_version: 1.8.6 # update to PTL version you saved ckpt on distributed_strategy: tensor_model_parallel_size: ${oc.decode:${oc.env:TP_DEGREE}} pipeline_model_parallel_size: 1 virtual_pipeline_model_parallel_size: 1 zero1: True sequence_parallel: True kv_replicator: ${oc.decode:${oc.env:KV_REPLICATOR}} data: micro_batch_size: 1 # limited by TRN memory global_batch_size: 8 train_dir: ${oc.env:DATA_DIR}/training.jsonl # [SFT/DPO] '/ubuntu/training.jsonl' or arrow file. As for ModelAlignment we use HF style dataloader, we also use HF style data file paths val_dir: ${oc.env:DATA_DIR}/validation.jsonl # [SFT/DPO] '/ubuntu/training.jsonl' or arrow file dev_choose_samples: null # [SFT] if set, will use those many number of records from the head of the dataset instead of using all. Set to null to use full dataset [SFT change 4 required] seq_length: 4096 # [SFT change 5 required] tokenizer: type: ${oc.env:ORIGINAL_MODEL_PATH} # [SFT change 6 required] model: # specify micro_batch_size, global_batch_size, and model parallelism # gradient accumulation will be done automatically based on data_parallel_size micro_batch_size: 1 # limited by TRN memory global_batch_size: 8 # will use more micro batches to reach global batch size # model architecture model_config: ${oc.env:MODEL_CONFIG_PATH} # TODO: Expand this into arguments in this file [SFT change 7 required] encoder_seq_length: 8192 max_position_embeddings: 8192 num_layers: 32 hidden_size: 4096 qkv_linear: True fuse_qkv: True rope_theta: 500000.0 transpose_nki_inputs: False # Miscellaneous use_cpu_initialization: True # Init weights on the CPU (slow for large models) weight_init_only: True # [SFT/DPO] Load only model states and ignore the optim states from ckpt directory ## Activation Checkpointing activations_checkpoint_granularity: selective # 'selective' or 'full' activations_checkpoint_recompute: [CoreAttention] fusions: softmax: True flash_attention: True do_layer_norm_weight_decay: True optim: name: adamw_fp32OptState lr: 1.5e-4 weight_decay: 0.01 capturable: False betas: - 0.9 - 0.999 sched: name: LinearAnnealingWithWarmUp warmup_steps: 10 max_steps: ${trainer.max_steps} model_alignment_strategy: sft: packing: True # [SFT] used for appending multiple records in a single record until seq length supported by model, if false uses pad tokens till seq length. True increases throughput precision: type: 'mixed_precision' # ['bf16SR', 'fp32', 'autocast', 'mixed_precision', 'mixed_precisionSR', 'manual'] # Set the following only if precision type is manual, otherwise they will be automatically set. master_weights: False fp32_grad_acc: False xla_use_bf16: '0' xla_downcast_bf16: '0' neuron_rt_stochastic_rounding_en: '0' compiler_flags: '--model-type transformer' compiler_cache_url: /home/ubuntu/neuron_cache aync_exec_max_inflight_requests: 5 bucket_size_collectives: 1024 neuron_rt_exec_timeout: 100 neuron_experimental_compress_rg: False💡 設定ファイルのFAQ
- 環境変数の値を設定ファイルに入れ込みたい
- 以下の記法で入れ込むことができます（これらの記法は yaml 自体の機能ではなく、yaml を解釈する OmegaConf ライブラリの機能です）
- そのまま文字列として入れる場合の例:
${oc.env:DATA_DIR} - 数値等に自動でパースさせる場合の例:
${oc.decode:${oc.env:KV_REPLICATOR}}
- そのまま文字列として入れる場合の例:
- 以下の記法で入れ込むことができます（これらの記法は yaml 自体の機能ではなく、yaml を解釈する OmegaConf ライブラリの機能です）
- 環境変数の値を設定ファイルに入れ込みたい
-
以下のコマンドで、
COMPILE=1環境変数付きで訓練スクリプトtrain.shを実行してください。cd ~/neuronx-distributed-training/examples COMPILE=1 sbatch --ntasks=1 --cpus-per-task=128 --wrap="srun ./train.sh"- 注意:
--cpus-per-task=128を指定し忘れると、処理が数十倍遅くなります。
- 注意:
-
-
AOTコンパイルのモニタリング：以下を確認してください。
- slurm のログファイル（デフォルトでは
slurm-<ジョブ名>.out）の中身を確認してください。ログの内容は、以下のようになります（重要な行のみを抜粋しています）-
最初に、jsonl 形式で与えられた学習データを、直ちに学習が可能なトークン列に変換する前処理が走ります
Map: 100%|██████████| 12008/12008 [00:12<00:00, 952.95 examples/s] -
その後、モデルの初期化が行われます
[NeMo I 2025-06-25 01:18:26 nlp_overrides:940] NLPTrainer: Initializing trainer with parameters: ...(省略) [2025-06-25 01:18:33.326: I neuronx_distributed/parallel_layers/parallel_state.py:592] > initializing tensor model parallel with size 32 [2025-06-25 01:18:33.501: I neuronx_distributed/parallel_layers/parallel_state.py:593] > initializing pipeline model parallel with size 1 [2025-06-25 01:18:33.501: I neuronx_distributed/parallel_layers/parallel_state.py:594] > initializing context model parallel with size 1 [2025-06-25 01:18:33.502: I neuronx_distributed/parallel_layers/parallel_state.py:595] > initializing data parallel with size 1 [2025-06-25 01:18:33.502: I neuronx_distributed/parallel_layers/parallel_state.py:596] > initializing world size to 32 [2025-06-25 01:22:08.856: I neuronx_distributed/parallel_layers/parallel_state.py:1390] [rank_0_pp0_tp0_dp0_cp0] > initializing kv group with size 4 -
その後、チェックポイントの読み込みが行われます（フルスクラッチウェイトからの学習の場合は除く）
[2025-06-30 08:47:34.316: I neuronx_distributed/trainer/checkpoint.py:887] loading checkpoint from pretrained_model [2025-06-30 08:48:28.400: I neuronx_distributed/trainer/checkpoint.py:963] loading checkpoint done -
その後、訓練ループが開始します。ただし、AOTコンパイルでは、実際の計算は行われず、計算グラフのキャプチャだけが行われます
Epoch 0: 0%| | 0/80 [00:00<?, ?it/s, reduced_train_loss=0.000, lr=0.000, parameter_norm=2.64e-38, global_step=0.000, consumed_samples=8.00, throughput=0.0249, throughput_peak=0.000]- AOTコンパイル中は実際の計算が行われていないため、ログに表示される
reduced_train_lossやparameter_normの値に意味はありません
- AOTコンパイル中は実際の計算が行われていないため、ログに表示される
-
ステップ数（グローバルステップ数）が 10 に到達すると、訓練ループを抜け、その間にキャプチャされた計算グラフたちをコンパイルする処理が開始します。コンパイラが起動されるたびに、以下のようなログが出力されます：
2025-06-25 04:05:43.000787: 75652 INFO ||NEURON_CC_WRAPPER||: Call compiler with cmd: neuronx-cc compile --framework=XLA /tmp/ubuntu/neuroncc_compile_workdir/4b8e6ed1-4425-4216-91e5-8f4a35c9768a/model.MODULE_10385560278789174898+4eb52b03.hlo_module.pb --output /tmp/ubuntu/neuroncc_compile_workdir/4b8e6ed1-4425-4216-91e5-8f4a35c9768a/model.MODULE_10385560278789174898+4eb52b03.neff --target=trn1 --model-type transformer --enable-mixed-precision-accumulation --verbose=35- yaml ファイル上の
trainer.max_stepsの設定にかかわらず、AOTコンパイルの際に走るステップ数は 10 となります（現状、これはtraining_orchestrator.pyにハードコードされています）。
- yaml ファイル上の
-
- 最終的に、ログが以下の両方を満たしている場合、AOTコンパイル成功です
-
最後に以下のような出力が表示されている（各種数値は状況によって変化しますが、Total failed compilations が 0 であることが必要です）
2025-06-25 05:58:16.000194: 58408 INFO ||NEURON_PARALLEL_COMPILE||: Total graphs: 18 2025-06-25 05:58:16.000194: 58408 INFO ||NEURON_PARALLEL_COMPILE||: Total successful compilations: 18 2025-06-25 05:58:16.000194: 58408 INFO ||NEURON_PARALLEL_COMPILE||: Total failed compilations: 0 -
文字列
ERRORを含むコンパイラ出力が含まれていない- 例えば以下のような行が途中で1つでも出力されていたらAOTコンパイルに失敗しています
2024-12-20 02:37:08.000427: 448565 ERROR ||NEURON_CC_WRAPPER||: Failed compilation with ...2025-Mar-27 17:23:42.108556 65653:71175 ERROR TDRV:tensor_allocate Failed to allocate 194641920 bytes on DEVICE for tensor UNKNOWN.2025-Mar-20 16:39:45.094874 158738:163246 ERROR ENC:ncclInitComm failed neuronInitComm request to NCCL2025-Jan-27 14:44:32.234379 199555:200462 ERROR NMGR:kmgr_load_nn_internal_v2 Failed to stage graph: kelf-0.json to NeuronCore
- 例えば以下のような行が途中で1つでも出力されていたらAOTコンパイルに失敗しています
-
- slurm のログファイル（デフォルトでは
-
学習の開始
-
AOTコンパイル時と同じコマンドを、
COMPILE=1環境変数無しで以下のように実行してください。sbatch --ntasks=1 --cpus-per-task=128 --wrap="srun ./train.sh"
-
-
学習のモニタリング：以下を確認してください。
- 学習の進行中に、コンピュートノードのうちの1つにssh接続して（手順は「slurm の基本的な使い方」を参照）、コンピュートノード上で neuron-top コマンドを実行してください。
- 32個全てのコアが使用されていることを確認してください。また、そのメモリ使用量を確認してください。
- slurmログの内容はAOTコンパイル時とほぼ同様ですが、今回は yaml ファイルに記述された
trainer.max_stepsのステップ数まで訓練ループが回ります。また、指定ステップ数（yaml 内のexp_manager.checkpoint_callback_params.every_n_train_steps）ごとにチェックポイントが保存されます。 - ログから以下を確認してください：
- 訓練1ステップが3秒ほどのオーダーであること
- チェックポイントファイルの保存間隔ステップ数が yaml での指定と合致していること
- チェックポイントファイルの保存に要する時間が10分程度のオーダーであること
- また、チェックポイントの保存内容を確認してください：
- チェックポイントの保存場所が yaml での指定（
exp_manager.exp_dirの配下）と合致していること - チェックポイントのファイルサイズを
duコマンドで確認し、model/配下が 21GB程度、optim/配下が91GB程度であること
- チェックポイントの保存場所が yaml での指定（
実はAOTコンパイルは必ずしも必要ではなく、いきなり学習（Just-in-Time コンパイル）を実行しても動作するケースが多いです。ただし、AOTコンパイル → 学習という二段階で実施するメリットとして、(A) JITコンパイルよりもAOTコンパイルの方が最適化が効きかつ合計コンパイル時間も短くなる、(B) 原因不明のエラーが出た際に問題の切り分けがしやすい、という点が挙げられます。
なお、一度学習を行った後、モデルアーキテクチャは固定のままで学習データだけ差し替えたような状況では、AOTコンパイルをし直す必要はありません（計算グラフは同じなので）。AOTコンパイルの結果は、コンパイラキャッシュディレクトリ（yaml 内の
compiler_cache_urlで指定）に吐き出されます。何らかの事情でキャッシュを削除したい場合は、下記の手順で削除が可能です（ただし再度AOTコンパイルが必要になります）。
コンパイルや学習実行中に意味不明なエラーが生じた場合、コンパイラキャッシュが悪さをしていることが（もしかすると）あるかもしれません。そのような場合、キャッシュを削除してみてください。- コンパイラキャッシュの削除方法について（公式Docs）
- 以下のコマンドでキャッシュを削除できます。
neuron_parallel_compile --command clean - コンパイルの予期せぬ中断によりロックファイルが残ってしまったものをクリアしたい場合は、以下のコマンドで可能です。
neuron_parallel_compile --command clear-locks - 設定ファイル上で、キャッシュディレクトリ (
compiler_cache_url) をデフォルトと異なるディレクトリに設定している場合は、上記の両方のコマンドを実行する前にあらかじめ環境変数NEURON_COMPILE_CACHE_URL=<cache URL>を設定しておく必要があります。
- 以下のコマンドでキャッシュを削除できます。
💡 訓練スクリプト train.sh についての FAQ
- 訓練ループはどこにありますか？
- 辿っていくとわかりますが、内部で
lightning.pytorch.trainerのTrainerを呼び出しており、この.fitの中にあります。
- 辿っていくとわかりますが、内部で
- 学習の進行中に、コンピュートノードのうちの1つにssh接続して（手順は「slurm の基本的な使い方」を参照）、コンピュートノード上で neuron-top コマンドを実行してください。
-
学習完了後のチェックポイントを NeuronX Distributed (NxD) 形式から HuggingFace Transformers (HF) 形式に変換します。
-
環境変数を設定します。
export OUTPUT_MODEL_DIR="/fsx/models/Tanuki-8B-dpo-v1.0-sft-step-100/"-
以下
convert_checkpoint_from_nxd_to_hf.shを保存した上で、sbatch convert_checkpoint_from_nxd_to_hf.shを実行してください（変換スクリプトの仕様上、変換はヘッドノードではなくコンピュートノード上で行う必要があります）。convert_checkpoint_from_nxd_to_hf.sh
#!/bin/bash echo "MODEL_CONFIG_PATH: $MODEL_CONFIG_PATH" echo "TP_DEGREE: $TP_DEGREE" echo "PP_DEGREE: $PP_DEGREE" echo "KV_REPLICATOR: $KV_REPLICATOR" cd ~/neuronx-distributed-training/examples/checkpoint_converter_scripts python3 checkpoint_converter.py \ --model_style hf \ --hw_backend trn1 \ --input_dir /fsx/experiments/hf_tanuki/28/checkpoints/hf_tanuki_8B_SFT--step=100-consumed_samples=728.0.ckpt/model \ --output_dir $OUTPUT_MODEL_DIR \ --load_xser True \ --config $MODEL_CONFIG_PATH \ --tp_size $TP_DEGREE \ --pp_size $PP_DEGREE \ --n_layers 32 \ --kv_size_multiplier $KV_REPLICATOR \ --qkv_linear True \ --fuse_qkv True \ --convert_to_full_state -
チャットテンプレートを書き換えて、トークナイザーを保存します。
temp="{% if messages[0]['role'] == 'system' %}{% set loop_messages = messages[1:] %}{% set system_message = messages[0]['content'] %}{% else %}{% set loop_messages = messages %}{% set system_message = false %}{% endif %}{% for message in loop_messages %}{% if (message['role'] == 'user') != (loop.index0 % 2 == 0) %}{{ raise_exception('Conversation roles must alternate user/assistant/user/assistant/...') }}{% endif %}{% if loop.index0 == 0 and system_message != false %}{% set content = '<<SYS>>\\n' + system_message + '\\n<</SYS>>\\n\\n' + message['content'] %}{% else %}{% set content = message['content'] %}{% endif %}{% if message['role'] == 'user' %}{{ bos_token + '[INST] ' + content.strip() + ' [/INST]' }}{% elif message['role'] == 'assistant' %}{{ ' ' + content.strip() + ' ' + eos_token }}{% endif %}{% endfor %}" curl https://huggingface.co/${HF_MODEL_NAME}/raw/main/tokenizer_config.json \ | jq --arg temp "${temp}" '. + {"chat_template": $temp}' \ > ${OUTPUT_MODEL_DIR}/tokenizer_config.json curl https://huggingface.co/${HF_MODEL_NAME}/raw/main/special_tokens_map.json -o ${OUTPUT_MODEL_DIR}/special_tokens_map.json curl https://huggingface.co/${HF_MODEL_NAME}/raw/main/tokenizer.json -o ${OUTPUT_MODEL_DIR}/tokenizer.json
-
- 最初のHF版の重みファイルと、学習完了後のHF版の重みファイルとで、パラメータの名前の一覧が完全に一致していることを確認してください。
- 最初のHF版の重みには、
model.safetensors.index.jsonファイルが含まれています。このファイルの中身を閲覧することで、パラメータ名の一覧を確認できます。 - 学習完了後のHF版の重みは、
checkpoint.pt形式の単一ファイルとなっています。checkpoint = torch.load("checkpoint.pt")により重み全体をメモリに読み込み、checkpoint.keys()を確認してください。
以上が、Trainium を使ってLLMを学習するひととおりの手順です。
今回の学習と全く同じモデルアーキテクチャ・モデル並列戦略で、同規模以下のサイズのモデルを学習する場合は、ここまでの手順を同様に踏めば問題ありません。しかし、それ以外のセッティングでの学習を行いたい場合、以降の知識が必要となります。
-
本50本ノックの内容を監修くださった AWS の常世様に、感謝を申し上げます。 ↩︎
Discussion