SageMaker HyperPod recipesを使ってfine-tuningを試してみた

$ aws --version

# aws-cli/2.17.56 Python/3.12.6 Linux/5.15.167.4-microsoft-standard-WSL2 exe/x86_64.ubuntu.24

$ curl 'https://static.us-east-1.prod.workshops.aws/public/97f372ae-83fb-43a1-9da2-0393b81dc89d/static/scripts/create_config.sh' --output create_config.sh

- : "${STACK_ID_VPC:=sagemaker-hyperpod}"
+ : "${STACK_ID_VPC:=cm-nakamura-hyperpod}"

$ export INSTANCES=g6e.12xlarge
$ bash create_config.sh 

export AWS_REGION=us-west-2
export INSTANCES=g6e.12xlarge
export VPC_ID=vpc-xxxx
export SUBNET_ID=subnet-xxxx
export PUBLIC_SUBNET_ID=subnet-xxxx
export FSX_ID=fs-xxxx
export FSX_MOUNTNAME=xxxx
export SECURITY_GROUP=sg-xxxx
export ROLE=arn:aws:iam::xxxx:role/cm-nakamura-hyperpod-AmazonSagemakerClusterExecutio-xxxx
export ROLENAME=cm-nakamura-hyperpod-AmazonSagemakerClusterExecutio-xxxx
export BUCKET=sagemaker-lifecycle-xxxx

$ source env_vars

$ git clone --depth=1 https://github.com/aws-samples/awsome-distributed-training/
$ cd awsome-distributed-training/1.architectures/5.sagemaker-hyperpod/LifecycleScripts/

# upload data
$ aws s3 cp --recursive base-config/ s3://${BUCKET}/src

# move back to env_var directory
$ cd ../../../..

$ aws s3 ls s3://${BUCKET}/src/

# 2025-01-12 23:11:26       1655 add_users.sh
# 2025-01-12 23:11:26        723 apply_hotfix.sh
# 2025-01-12 23:11:26       3119 config.py
# 2025-01-12 23:11:26       8948 lifecycle_script.py
# 2025-01-12 23:11:26       3226 mount_fsx.sh
# 2025-01-12 23:11:26       1337 on_create.sh
# 2025-01-12 23:11:26       3577 setup_mariadb_accounting.sh
# 2025-01-12 23:11:26       2843 setup_rds_accounting.sh
# 2025-01-12 23:11:27       8262 setup_sssd.py
# 2025-01-12 23:11:27         88 shared_users_sample.txt
# 2025-01-12 23:11:27       1355 start_slurm.sh

$ source env_vars
$ cat > cluster-config.json << EOL
{
    "ClusterName": "cm-nakamura-cluster",
    "InstanceGroups": [
      {
        "InstanceGroupName": "controller-machine",
        "InstanceType": "ml.c5.xlarge",
        "InstanceStorageConfigs": [
          {
            "EbsVolumeConfig": {
              "VolumeSizeInGB": 500
            }
          }
        ],
        "InstanceCount": 1,
        "LifeCycleConfig": {
          "SourceS3Uri": "s3://${BUCKET}/src",
          "OnCreate": "on_create.sh"
        },
        "ExecutionRole": "${ROLE}",
        "ThreadsPerCore": 2
      },
      {
        "InstanceGroupName": "worker-group-1",
        "InstanceType": "ml.g6e.12xlarge",
        "InstanceCount": 1,
        "InstanceStorageConfigs": [
          {
            "EbsVolumeConfig": {
              "VolumeSizeInGB": 500
            }
          }
        ],
        "LifeCycleConfig": {
          "SourceS3Uri": "s3://${BUCKET}/src",
          "OnCreate": "on_create.sh"
        },
        "ExecutionRole": "${ROLE}",
        "ThreadsPerCore": 1
      }
    ],
    "VpcConfig": {
      "SecurityGroupIds": ["$SECURITY_GROUP"],
      "Subnets":["$SUBNET_ID"]
    }
}
EOL

$ instance_type=$(jq '.InstanceGroups[] | select(.InstanceGroupName == "worker-group-1").InstanceType' cluster-config.json)
$ cat > provisioning_parameters.json << EOL
{
  "version": "1.0.0",
  "workload_manager": "slurm",
  "controller_group": "controller-machine",
  "worker_groups": [
    {
      "instance_group_name": "worker-group-1",
      "partition_name": ${instance_type}
    }
  ],
  "fsx_dns_name": "${FSX_ID}.fsx.${AWS_REGION}.amazonaws.com",
  "fsx_mountname": "${FSX_MOUNTNAME}"
}
EOL

# copy to the S3 Bucket
$ aws s3 cp provisioning_parameters.json s3://${BUCKET}/src/

$ aws sagemaker create-cluster \
    --cli-input-json file://cluster-config.json \
    --region $AWS_REGION

$ aws sagemaker list-clusters --output table

# ---------------------------------------------------------------------------------------------------------------------------------------------------------
# |                                                                     ListClusters                                                                      |
# +-------------------------------------------------------------------------------------------------------------------------------------------------------+
# ||                                                                  ClusterSummaries                                                                   ||
# |+----------------------------------------------------------------+------------------------------+----------------+------------------------------------+|
# ||                           ClusterArn                           |         ClusterName          | ClusterStatus  |           CreationTime             ||
# |+----------------------------------------------------------------+------------------------------+----------------+------------------------------------+|
# ||  arn:aws:sagemaker:us-west-2:001234567899:cluster/{CLUSTER_ID} |  cm-nakamura-cluster         |  Creating      |  2025-01-13T11:21:53.178000+09:00  ||
# |+----------------------------------------------------------------+------------------------------+----------------+------------------------------------+|

$ sudo curl "https://s3.amazonaws.com/session-manager-downloads/plugin/latest/ubuntu_64bit/session-manager-plugin.deb" -o "/tmp/session-manager-plugin.deb"
$ sudo dpkg -i /tmp/session-manager-plugin.deb

$ cp ~/.ssh/config ~/.ssh/config.bak

$ ssh-keygen -t rsa -q -f "$HOME/.ssh/hyperpod" -N ""

$ ls -l $HOME/.ssh/ | grep "hyperpod"

# -rw------- 1 nakamura nakamura 2602 Jan 13 11:59 hyperpod
# -rw-r--r-- 1 nakamura nakamura  570 Jan 13 11:59 hyperpod.pub

$ curl -O https://raw.githubusercontent.com/aws-samples/awsome-distributed-training/main/1.architectures/5.sagemaker-hyperpod/easy-ssh.sh

# Function to add the user's SSH public key to the cluster
add_keypair_to_cluster() {
-     PUBLIC_KEY=$(cat ~/.ssh/id_rsa.pub)
+     PUBLIC_KEY=$(cat ~/.ssh/hyperpod.pub)

    # Check if the fingerprint already exists in the cluster's authorized_keys
    EXISTING_KEYS=$(aws ssm start-session --target sagemaker-cluster:${cluster_id}_${node_group}-${instance_id} --document-name AmazonEKS-ExecuteNonInteractiveCommand --parameters command="cat /fsx/ubuntu/.ssh/authorized_keys")

    if echo "$EXISTING_KEYS" | grep -q "$PUBLIC_KEY"; then
-         echo -e "${BLUE}2. Detected SSH public key ${GREEN}~/.ssh/id_rsa.pub${BLUE} on the cluster. Skipping adding...${NC}" 
+         echo -e "${BLUE}2. Detected SSH public key ${GREEN}~/.ssh/hyperpod.pub${BLUE} on the cluster. Skipping adding...${NC}" 
        return
    else
-         echo -e "${BLUE}2. Do you want to add your SSH public key ${GREEN}~/.ssh/id_rsa.pub${BLUE} to the cluster (yes/no)?${NC}" 
+         echo -e "${BLUE}2. Do you want to add your SSH public key ${GREEN}~/.ssh/hyperpod.pub${BLUE} to the cluster (yes/no)?${NC}" 
        read -p "> " ADD_KEYPAIR
        if [[ $ADD_KEYPAIR == "yes" ]]; then
            echo "Adding ... ${PUBLIC_KEY}"
            command="sed -i \$a$(escape_spaces "$PUBLIC_KEY") /fsx/ubuntu/.ssh/authorized_keys"
            aws ssm start-session --target sagemaker-cluster:${cluster_id}_${node_group}-${instance_id}  --document-name AmazonEKS-ExecuteNonInteractiveCommand  --parameters command="$command"
-             echo "✅ Your SSH public key ~/.ssh/id_rsa.pub has been added to the cluster."
+             echo "✅ Your SSH public key ~/.ssh/hyperpod.pub has been added to the cluster."
        else
            echo "❌ Skipping adding SSH public key to the cluster."
        fi
    fi
}

$ ./easy-ssh.sh -c controller-machine cm-nakamura-cluster

# =================================================
# 
# ==== 🚀 HyperPod Cluster Easy SSH Script! 🚀 ====
# 
# 
# =================================================
# Cluster id: xxxx
# Instance id: xxxx
# Node Group: controller-machine
# Would you like to add cm-nakamura-cluster to  ~/.ssh/config (yes/no)?
# > yes
# ✅ adding ml-cluster to  ~/.ssh/config:
# 2. Do you want to add your SSH public key ~/.ssh/hyperpod.pub to the cluster (yes/no)?
# > yes
# Adding ... ssh-rsa xxxx
# 
# Starting session with SessionId: cm-nakamura.shogo-xxxx
# 
# 
# Exiting session with sessionId: cm-nakamura.shogo-xxxx.
# 
# ✅ Your SSH public key ~/.ssh/hyperpod.pub has been added to the cluster.
# 
# Now you can run:
# 
# $ ssh cm-nakamura-cluster
# 
# Starting session with SessionId: cm-nakamura.shogo-xxxx
# /bin/bash -c 'cd ~ && exec bash'

ubuntu@ip-xxxx:~$ 

Host cm-nakamura-cluster
    User ubuntu
+     IdentityFile ~/.ssh/hyperpod
    ProxyCommand sh -c "aws ssm start-session  --target sagemaker-cluster:xxxx_controller-machine-xxxx --document-name AWS-StartSSHSession --parameters 'portNumber=%p'"

$ ssh cm-nakamura-cluster

# ========================================================================================
#     ____              __  ___     __             __ __                  ___          __
#    / __/__ ____ ____ /  |/  /__ _/ /_____ ____  / // /_ _____  ___ ____/ _ \___  ___/ /
#   _\ \/ _ `/ _ `/ -_) /|_/ / _ `/  '_/ -_) __/ / _  / // / _ \/ -_) __/ ___/ _ \/ _  /
#  /___/\_,_/\_, /\__/_/  /_/\_,_/_/\_\\__/_/   /_//_/\_, / .__/\__/_/ /_/   \___/\_,_/
#           /___/                                    /___/_/
#                           HyperPod Instance AMI (Ubuntu 20.04)
# ========================================================================================

# Welcome to Ubuntu 20.04.6 LTS (GNU/Linux 5.15.0-1072-aws x86_64v)

# Utility libraries are installed in /usr/bin/python3.9.
# To access them, use /usr/bin/python3.9.

# AWS Deep Learning AMI Homepage: https://aws.amazon.com/machine-learning/amis/
# Release Notes: https://docs.aws.amazon.com/dlami/latest/devguide/appendix-ami-release-notes.html
# Support: https://forums.aws.amazon.com/forum.jspa?forumID=263
# For a fully managed experience, check out Amazon SageMaker at https://aws.amazon.com/sagemaker
# =============================================================================

#  * Documentation:  https://help.ubuntu.com
#  * Management:     https://landscape.canonical.com
#  * Support:        https://ubuntu.com/pro

#  System information as of Mon Jan 13 03:11:39 UTC 2025

#   System load:  0.0                Processes:             245
#   Usage of /:   52.2% of 96.73GB   Users logged in:       0
#   Memory usage: 14%                IPv4 address for ens6: xxxx
#   Swap usage:   0%

#  * Ubuntu Pro delivers the most comprehensive open source security and
#    compliance features.

#    https://ubuntu.com/aws/pro

# Expanded Security Maintenance for Applications is not enabled.

# 19 updates can be applied immediately.
# 11 of these updates are standard security updates.
# To see these additional updates run: apt list --upgradable

# 41 additional security updates can be applied with ESM Apps.
# Learn more about enabling ESM Apps service at https://ubuntu.com/esm

# New release '22.04.5 LTS' available.
# Run 'do-release-upgrade' to upgrade to it.

# To replace an instance run:
#    sudo scontrol update node=<hostname> state=fail reason="Action:Replace"

# To automatically resume jobs, please add the following in your job submission script:
#    srun --auto-resume=1

# Instance Type: c5.xlarge

# =============================================================================
# AMI Name: Deep Learning Base OSS Nvidia Driver GPU AMI (Ubuntu 20.04) 
# Supported EC2 instances: G4dn, G5, G6, Gr6, G6e, P4d, P4de, P5, P5e, P5en, Trn1, Trn1n
# NVIDIA driver version: 550.127.05
# CUDA versions available: cuda-12.1 cuda-12.2 cuda-12.3 cuda-12.4
# Default CUDA version is 12.1

# Release notes: https://docs.aws.amazon.com/dlami/latest/devguide/appendix-ami-release-notes.html
# AWS Deep Learning AMI Homepage: https://aws.amazon.com/machine-learning/amis/
# Developer Guide and Release Notes: https://docs.aws.amazon.com/dlami/latest/devguide/what-is-dlami.html
# Support: https://forums.aws.amazon.com/forum.jspa?forumID=263
# For a fully managed experience, check out Amazon SageMaker at https://aws.amazon.com/sagemaker
# =============================================================================
#    ____              __  ___     __             __ __                  ___          __
#   / __/__ ____ ____ /  |/  /__ _/ /_____ ____  / // /_ _____  ___ ____/ _ \___  ___/ /
#  _\ \/ _ `/ _ `/ -_) /|_/ / _ `/  '_/ -_) __/ / _  / // / _ \/ -_) __/ ___/ _ \/ _  /
# /___/\_,_/\_, /\__/_/  /_/\_,_/_/\_\\__/_/   /_//_/\_, / .__/\__/_/ /_/   \___/\_,_/
#          /___/                                    /___/_/

# To replace an instance run:
#    sudo scontrol update node=<hostname> state=fail reason="Action:Replace"

# To automatically resume jobs, please add the following in your job submission script:
#    srun --auto-resume=1

# You're on the login
# Controller Node IP: xxxx
# Login Node IP: xxxx
# Instance Type: ml.c5.xlarge
# Last login: Mon Jan 13 03:

ubuntu@ip-xxxx:~$ 

$ ./easy-ssh.sh -c worker-group-1 cm-nakamura-cluster

$ export DATASET_BUCKET_NAME="cm-nakamura-dataset"

$ aws s3api create-bucket --bucket ${DATASET_BUCKET_NAME} --region ${AWS_REGION} \
    --create-bucket-configuration LocationConstraint=${AWS_REGION}

$ aws fsx create-data-repository-association \
    --region ${AWS_REGION} \
    --file-system-id ${FSX_ID} \
    --file-system-path /dataset \
    --data-repository-path s3://${DATASET_BUCKET_NAME} \
    --s3 "AutoImportPolicy={Events=[NEW,CHANGED,DELETED]},AutoExportPolicy={Events=[NEW,CHANGED,DELETED]}" \
    --batch-import-meta-data-on-create

$ aws s3 cp ./tokenized_dataset.tar s3://${DATASET_BUCKET_NAME}

$ cd /fsx/dataset
$ sudo tar xvf tokenized_dataset.tar

$ tree

# .
# ├── tokenized_dataset
# │   ├── test
# │   │   ├── data-00000-of-00001.arrow
# │   │   ├── dataset_info.json
# │   │   └── state.json
# │   ├── train
# │   │   ├── data-00000-of-00008.arrow
# │   │   ├── data-00001-of-00008.arrow
# │   │   ├── data-00002-of-00008.arrow
# │   │   ├── data-00003-of-00008.arrow
# │   │   ├── data-00004-of-00008.arrow
# │   │   ├── data-00005-of-00008.arrow
# │   │   ├── data-00006-of-00008.arrow
# │   │   ├── data-00007-of-00008.arrow
# │   │   ├── dataset_info.json
# │   │   └── state.json
# │   └── validation
# │       ├── data-00000-of-00001.arrow
# │       ├── dataset_info.json
# │       └── state.json
# └── tokenized_dataset.tar

$ sudo apt install python3.9-venv
$ sudo apt install python3.9-dev

$ python3.9 -m venv ${PWD}/venv
$ source venv/bin/activate

(venv) $ git clone --recursive https://github.com/aws/sagemaker-hyperpod-recipes.git
(venv) $ cd sagemaker-hyperpod-recipes
(venv) $ pip3 install --upgrade setuptools wheel
(venv) $ pip3 install -r requirements.txt

$ export AWS_REGION="us-west-2"
$ REGION=$AWS_REGION
$ IMAGE="658645717510.dkr.ecr.${REGION}.amazonaws.com/smdistributed-modelparallel:2.4.1-gpu-py311-cu121"
$ aws ecr get-login-password --region "${REGION}" | docker login --username AWS --password-stdin 658645717510.dkr.ecr.${REGION}.amazonaws.com

# Login Succeeded

$ enroot import -o $PWD/smdistributed-modelparallel.sqsh dockerd://${IMAGE}

# [INFO] Fetching image
# 
# Unable to find image '658645717510.dkr.ecr.us-west-2.amazonaws.com/smdistributed-modelparallel:2.4.1-gpu-py311-cu121' locally
# 2.4.1-gpu-py311-cu121: Pulling from smdistributed-modelparallel
# 7a2c55901189: Pulling fs layer 
# 755e535b54a3: Pulling fs layer 
#
# ...(中略)...
#
# Digest: sha256:a0e2cefb276671c79728b7f236ec50570535400843a4db5e9715930e3ae1c007
# Status: Downloaded newer image for 658645717510.dkr.ecr.us-west-2.amazonaws.com/smdistributed-modelparallel:2.4.1-gpu-py311-cu121
# 95037b48fcd3d6e025cb6a44a22c0afb3518d6c4b496e6688a2c1baeedfaf585
# 
# [INFO] Extracting image content...
# [INFO] Creating squashfs filesystem...
# 
# Parallel mksquashfs: Using 4 processors
# Creating 4.0 filesystem on /fsx/ubuntu/sagemaker-hyperpod-recipes/smdistributed-modelparallel.sqsh, block size 131072.
# [===================================================================================================\] 346971/346971 100%
# 
# Exportable Squashfs 4.0 filesystem, lzo compressed, data block size 131072
#         uncompressed data, uncompressed metadata, uncompressed fragments,
#         uncompressed xattrs, uncompressed ids
#         duplicates are not removed
# Filesystem size 19724778.70 Kbytes (19262.48 Mbytes)
#         99.87% of uncompressed filesystem size (19750283.71 Kbytes)
# Inode table size 8522800 bytes (8323.05 Kbytes)
#         100.00% of uncompressed inode table size (8522800 bytes)
# Directory table size 6499860 bytes (6347.52 Kbytes)
#         100.00% of uncompressed directory table size (6499860 bytes)
# No duplicate files removed
# Number of inodes 242167
# Number of files 207963
# Number of fragments 14355
# Number of symbolic links  10322
# Number of device nodes 0
# Number of fifo nodes 0
# Number of socket nodes 0
# Number of directories 23882
# Number of ids (unique uids + gids) 1
# Number of uids 1
#         root (0)
# Number of gids 1
#         root (0)

#!/bin/bash

# Original Copyright (c), NVIDIA CORPORATION. Modifications © Amazon.com

#Users should setup their cluster type in /recipes_collection/config.yaml

SAGEMAKER_TRAINING_LAUNCHER_DIR=${SAGEMAKER_TRAINING_LAUNCHER_DIR:-"$(pwd)"}

- HF_MODEL_NAME_OR_PATH="${HF_MODEL_NAME_OR_PATH}" # HuggingFace pretrained model name or path
- HF_ACCESS_TOKEN="${HF_ACCESS_TOKEN}" # Optional HuggingFace access token
+ HF_MODEL_NAME_OR_PATH="meta-llama/Meta-Llama-3-8B" # HuggingFace pretrained model name or path
+ HF_ACCESS_TOKEN="{ここにHugging Faceのトークンを記載}" # Optional HuggingFace access token

- TRAIN_DIR="${TRAIN_DIR}" # Location of training dataset
- VAL_DIR="${VAL_DIR}" # Location of validation dataset
+ TRAIN_DIR="/fsx/dataset/tokenized_dataset/train" # Location of training dataset
+ VAL_DIR="/fsx/dataset/tokenized_dataset/validation" # Location of validation dataset

- EXP_DIR="${EXP_DIR}" # Location to save experiment info including logging, checkpoints, etc.
+ EXP_DIR="/fsx/ubuntu/exp" # Location to save experiment info including logging, checkpoints, etc.

HYDRA_FULL_ERROR=1 python3 "${SAGEMAKER_TRAINING_LAUNCHER_DIR}/main.py" \
    base_results_dir="${SAGEMAKER_TRAINING_LAUNCHER_DIR}/results" \
+     instance_type=ml.g6e.12xlarge \
+     container=/fsx/ubuntu/smdistributed-modelparallel.sqsh \
    recipes=fine-tuning/llama/p4_hf_llama3_8b_seq8k_gpu_lora \
    recipes.run.name="hf-llama3-8b-lora" \
    recipes.exp_manager.exp_dir="$EXP_DIR" \
    recipes.trainer.num_nodes=1 \
    recipes.model.train_batch_size=1 \
    recipes.model.data.train_dir="$TRAIN_DIR" \
    recipes.model.data.val_dir="$VAL_DIR" \
    recipes.model.hf_model_name_or_path="$HF_MODEL_NAME_OR_PATH" \
    recipes.model.hf_access_token="$HF_ACCESS_TOKEN" \
+     recipes.model.shard_degree=4 \
+     recipes.trainer.devices=4 \

- HF_MODEL_NAME_OR_PATH="${HF_MODEL_NAME_OR_PATH}" # HuggingFace pretrained model name or path
- HF_ACCESS_TOKEN="${HF_ACCESS_TOKEN}" # Optional HuggingFace access token
+ HF_MODEL_NAME_OR_PATH="meta-llama/Meta-Llama-3-8B" # HuggingFace pretrained model name or path
+ HF_ACCESS_TOKEN="{ここにHugging Faceのトークンを記載}" # Optional HuggingFace access token

- TRAIN_DIR="${TRAIN_DIR}" # Location of training dataset
- VAL_DIR="${VAL_DIR}" # Location of validation dataset
+ TRAIN_DIR="/fsx/dataset/tokenized_dataset/train" # Location of training dataset
+ VAL_DIR="/fsx/dataset/tokenized_dataset/validation" # Location of validation dataset

- EXP_DIR="${EXP_DIR}" # Location to save experiment info including logging, checkpoints, etc.
+ EXP_DIR="/fsx/ubuntu/exp" # Location to save experiment info including logging, checkpoints, etc.

+     instance_type=ml.g6e.12xlarge \

+     container=/fsx/ubuntu/smdistributed-modelparallel.sqsh \

+     recipes.trainer.num_nodes=1 \
+     recipes.model.shard_degree=4 \
+     recipes.trainer.devices=4 \

# Original Copyright (c), NVIDIA CORPORATION. Modifications © Amazon.com

exclusive: True
mem: 0
job_name_prefix: 'sagemaker-'
slurm_create_submission_file_only: False # Setting to True if just want to create submission file
stderr_to_stdout: True # Setting to False to split the stderr and stdout logs
srun_args:
  # - "--no-container-mount-home"
slurm_docker_cfg:
  docker_args:
    # - "--runtime=nvidia" # this is required if the docker runtime version is low
  post_launch_commands: # commands will run after launching the docker container using bash
- container_mounts: # List of additional paths to mount to container. They will be mounted to same path.
-   - null
+ container_mounts:
+   - /fsx/dataset:/fsx/dataset
+   - /fsx/ubuntu/exp:/fsx/ubuntu/exp

$ bash ./launcher_scripts/llama/p4_run_hf_llama3_8b_seq8k_gpu_lora.sh

$ squeue

#              JOBID PARTITION     NAME     USER ST       TIME  NODES NODELIST(REASON)
#                  3       dev sagemake   ubuntu  R       4:06      1 ip-10-1-57-144

$ sinfo

# PARTITION       AVAIL  TIMELIMIT  NODES  STATE NODELIST
# dev*               up   infinite      1  alloc ip-10-1-57-144
# ml.g6e.12xlarge    up   infinite      1  alloc ip-10-1-57-144

$ tree

# /fsx/ubuntu/exp
# ├── checkpoints
# │   └── peft_full
# │       └── steps_50
# │           ├── README.md
# │           ├── adapter_config.json
# │           ├── adapter_model.safetensors
# │           └── final-model
# │               ├── config.json
# │               ├── generation_config.json
# │               ├── model-00001-of-00004.safetensors
# │               ├── model-00002-of-00004.safetensors
# │               ├── model-00003-of-00004.safetensors
# │               ├── model-00004-of-00004.safetensors
# │               └── model.safetensors.index.json
# └── experiment
#     ├── 2025-01-22_04-32-34
#     │   ├── cmd-args.log
#     │   ├── events.out.tfevents.1737520360.ip-10-1-78-125.41029.0
#     │   ├── git-info.log
#     │   ├── hparams.yaml
#     │   ├── lightning_logs.txt
#     │   ├── nemo_error_log.txt
#     │   └── sagemaker_log_globalrank-0_localrank-0.txt
#     ├── sagemaker_log_globalrank-1_localrank-1.txt
#     ├── sagemaker_log_globalrank-2_localrank-2.txt
#     └── sagemaker_log_globalrank-3_localrank-3.txt