GKE GPU A3 with TCPX support

[samos123]

Add the ability to launch GPU jobs directly on GKE

Test cases I ran:

Fuji 7b single host per device batch 8 is working too:

BASTION_TIER=1 axlearn gcp gke start --instance_type=gpu-a3-highgpu-8g-8 \
        --cluster=daily-run-single-zone-gker27 --num_replicas=1 \
        --queue multislice-queue \
        --bundler_spec=allow_dirty=True \
        --bundler_type=artifactregistry --bundler_spec=image=gpu \
        --bundler_spec=dockerfile=Dockerfile --bundler_spec=target=gpu \
        -- mkdir -p /tmp/test_trainer\; python3 -m axlearn.common.launch_trainer_main \
  --module=text.gpt.c4_trainer --config=fuji-7B-v1-single-host \
  --trainer_dir=/tmp/test_trainer --data_dir=gs://axlearn-public/tensorflow_datasets \
  --jax_backend=gpu

step time on single node: 2.06 seconds

Fuji 7b 32 nodes batch per device 8:

BASTION_TIER=1 axlearn gcp gke start --instance_type=gpu-a3-highgpu-8g-256 \
        --cluster=daily-run-single-zone-gker27 --num_replicas=32 \
        --queue multislice-queue \
        --bundler_spec=allow_dirty=True \
        --bundler_type=artifactregistry --bundler_spec=image=gpu \
        --bundler_spec=dockerfile=Dockerfile --bundler_spec=target=gpu \
        -- mkdir -p /tmp/test_trainer\; python3 -m axlearn.common.launch_trainer_main \
  --module=text.gpt.c4_trainer --config=fuji-7B-v1 \
  --trainer_dir=/tmp/test_trainer --data_dir=gs://axlearn-public/tensorflow_datasets \
  --jax_backend=gpu

step time on 32 nodes: 3.03 seconds

Fuji 7b 64 nodes batch per device 4:

BASTION_TIER=1 axlearn gcp gke start --instance_type=gpu-a3-highgpu-8g-512 \
        --cluster=daily-run-single-zone-gker27 --num_replicas=64 \
        --queue multislice-queue \
        --bundler_spec=allow_dirty=True \
        --bundler_type=artifactregistry --bundler_spec=image=gpu \
        --bundler_spec=dockerfile=Dockerfile --bundler_spec=target=gpu \
        -- mkdir -p /tmp/test_trainer\; python3 -m axlearn.common.launch_trainer_main \
  --module=text.gpt.c4_trainer --config=fuji-7B-v1 \
  --trainer_dir=/tmp/test_trainer --data_dir=gs://axlearn-public/tensorflow_datasets \
  --jax_backend=gpu

Step time A3 3.92 seconds (note A3+ this is 1.785 seconds).

I plan to work on A3+ support in a separate PR.

[markblee]

A quick skim through...

[markblee]

Thanks, overall looks reasonable. Can we add/update tests for changes in job.py and gke_runner.py?

[samos123]

I have addressed all comments, some with fixes and some with responses instead.

I've added tests for GPUGKEJob. I also added some basic tests in the gke_runner

[markblee]

Thanks! Looks good to me after final comments.

[samos123]

I did another test on single node but still waiting for multi-node results. The single node result was able to verify that ; works instead of \n.

[markblee]

Can we leave a comment "Done" for comments which have been addressed? (See also: https://github.com/apple/axlearn/blob/main/CONTRIBUTING.md)

[samos123]

I resolved all comments and fixed missing punctation

[markblee]

Thanks!

[samos123]

Seems Circle CI is awaiting someone to approve the run?

[samos123]

Seems it was approved but there was a Circe CI Apple Auth error: https://app.circleci.com/pipelines/github/apple/axlearn/1925/workflows/99dbcc0a-7a9f-4338-8ec3-0cc2724a9e28
I can't see what the error is though. So will probably need @markblee to take a look?

[samos123]

Looks like Apple Authenticate in circle ci is still failing but I don't see the logs.