Skip to content

Conversation

@lkolluru05
Copy link
Contributor

@lkolluru05 lkolluru05 commented Oct 7, 2025

`#### Different names for these images because they are in same repository #####

export NAME=axlearn-img
export COLOCATED_NAME=colocated-img
export CKPT_BUCKET_NAME=<>

axlearn gcp bundle --name=$NAME
--bundler_spec=allow_dirty=True
--bundler_type=artifactregistry
--bundler_spec=dockerfile=Dockerfile
--bundler_spec=image=tpu
--bundler_spec=target=tpu
--bundler_spec=colocated_image_required=True
--bundler_spec=colocated_image_name=$COLOCATED_NAME

axlearn gcp launch run --cluster=mlperf-v5p
--runner_name gke_tpu_pathways \
--name=$NAME \
--instance_type=tpu-v5p-32
--num_replicas=1 \
--bundler_spec=allow_dirty=True \
--bundler_type=artifactregistry
--bundler_spec=image=tpu \
--bundler_spec=dockerfile=Dockerfile
--bundler_spec=target=tpu
--colocated_image=$COLOCATED_NAME \
-- TPU_PREMAPPED_BUFFER_SIZE=34359738368 python3 test_benchmark.py --ckpt_path $CKPT_BUCKET_NAME
`

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants