distributed-training-trn1_custom_ami.yaml

# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: Apache-2.0

# For additional examples please refer to this [Github repository](https://github.com/aws-neuron/aws-neuron-parallelcluster-samples/blob/master/examples/jobs/neuronx-nemo-megatron-llamav2-job.md) from aws-neuron.

Imds:
  ImdsSupport: v2.0
Image:
  Os: alinux2
  CustomAmi: PLACEHOLDER_CUSTOM_AMI_ID
HeadNode:
  InstanceType: m5.8xlarge
  Networking:
    SubnetId: PLACEHOLDER_PUBLIC_SUBNET
  Ssh:
    KeyName: PLACEHOLDER_SSH_KEY
  LocalStorage:
    RootVolume:
      Size: 500
      DeleteOnTermination: true # that's your root and /home volume for users
  Iam:
    AdditionalIamPolicies: # grant ECR, SSM and S3 read access
      - Policy: arn:aws:iam::aws:policy/AmazonSSMManagedInstanceCore
      - Policy: arn:aws:iam::aws:policy/AmazonS3ReadOnlyAccess
      - Policy: arn:aws:iam::aws:policy/AmazonEC2ContainerRegistryReadOnly
Scheduling:
  Scheduler: slurm
  SlurmSettings:
    ScaledownIdletime: 60
  SlurmQueues:
    - Name: compute-gpu
      CapacityType: ONDEMAND
      Networking:
        SubnetIds:
          - PLACEHOLDER_PRIVATE_SUBNET
      ComputeSettings:
        LocalStorage:
          EphemeralVolume:
            MountDir: /local_scratch # each instance has a local scratch on NVMe
          RootVolume:
            Size: 200
      # The capacity reservation section is recommended if you use instances
      # with a targeted ODCRs. You can also use a capacity resource group and
      # CapacityReservationResourceGroupArn if you want to regroup
      # multiple reservations
      ComputeResources:
        - Name: distributed-ml
          InstanceType: trn1.32xlarge
          MinCount: PLACEHOLDER_MIN_INSTANCES # if min = max then capacity is maintained and will
          MaxCount: PLACEHOLDER_MAX_INSTANCES # not scale down
          Efa:
            Enabled: true
          # assumes you are using a capacity reservation.
          # If not comment or remove the 2 lines below
          CapacityReservationTarget:
            CapacityReservationId: PLACEHOLDER_CAPACITY_RESERVATION_ID
          Networking:
            PlacementGroup:
              Enabled: true
              Id: PLACEHOLDER_PLACEMENT_GROUP
SharedStorage:
  - MountDir: /fsx
    Name: fsx
    StorageType: FsxLustre
    FsxLustreSettings:
      StorageCapacity: 4800 # size it to your storage and throughput needs
      PerUnitStorageThroughput: 250 # this can be increased to 500 and 100
      DeploymentType: PERSISTENT_2
  - Name: SharedEBS
    StorageType: Ebs
    MountDir: /apps # Store your shared apps & scripts here
    EbsSettings:
      VolumeType: gp3
      Size: 200
      Throughput: 300
      Iops: 6000
Monitoring:
  DetailedMonitoring: true
  Logs:
    CloudWatch:
      Enabled: true # good for debug
  Dashboards:
    CloudWatch:
      Enabled: false # provide basic dashboards