accelerate_job.json

{
    "displayName": "brittrock_hf_laion_bs1_16_real",
    "jobSpec": {
        "workerPoolSpecs": [
            {
                "machineSpec": {
                    "machineType": "a3-megagpu-8g",
                    "acceleratorType": "NVIDIA_H100_MEGA_80GB",
                    "acceleratorCount": 8,
                    "reservationAffinity": {
                        "reservationAffinityType": "SPECIFIC_RESERVATION",
                        "key": "compute.googleapis.com/reservation-name",
                        "values": [
                            "projects/disco-sector-292704/zones/us-east4-a/reservations/snap-a3-mega-reservation"
                        ]
                    }
                },
                "replicaCount": "1",
                "containerSpec": {
                    "imageUri": "us-east5-docker.pkg.dev/google.com/vertex-training-dlexamples/nemo-sd-training-repository/sd-accelerate_train:latest",
                    "command": [
                        "sh",
                        "-c"
                    ],
                    "args": [
                        "git clone https://github.com/bvrockwell/hf-multi.git && chmod +x hf-multi/set_env.sh && ./hf-multi/set_env.sh"
                    ],
                    "env": [
                        {
                            "name": "LD_LIBRARY_PATH",
                            "value": "/usr/local/nvidia/lib64"
                        },
                        {
                            "name": "NODE_COUNT",
                            "value": "2"
                        },
                        {
                            "name": "ACC_CONFIG",
                            "value": "/hf-multi/2host_config.yaml"
                        }
                    ]
                }
            },
            {
                "machineSpec": {
                    "machineType": "a3-megagpu-8g",
                    "acceleratorType": "NVIDIA_H100_MEGA_80GB",
                    "acceleratorCount": 8,
                    "reservationAffinity": {
                        "reservationAffinityType": "SPECIFIC_RESERVATION",
                        "key": "compute.googleapis.com/reservation-name",
                        "values": [
                            "projects/disco-sector-292704/zones/us-east4-a/reservations/snap-a3-mega-reservation"
                        ]
                    }
                },
                "replicaCount": "1",
                "containerSpec": {
                    "imageUri": "us-east5-docker.pkg.dev/google.com/vertex-training-dlexamples/nemo-sd-training-repository/sd-accelerate_train:latest",
                    "command": [
                        "sh",
                        "-c"
                    ],
                    "args": [
                        "git clone https://github.com/bvrockwell/hf-multi.git && chmod +x hf-multi/set_env.sh && ./hf-multi/set_env.sh"
                    ],
                    "env": [
                        {
                            "name": "LD_LIBRARY_PATH",
                            "value": "/usr/local/nvidia/lib64"
                        },
                        {
                            "name": "NODE_COUNT",
                            "value": "2"
                        },
                        {
                            "name": "ACC_CONFIG",
                            "value": "/hf-multi/2host_config.yaml"
                        }
                    ]
                }
            }
        ],
        "enableWebAccess": true
    }
}