AXLearn

AXLearn is a deep learning design framework, built on top of JAX and XLA, to support the development of large-scale models.

Hardware and Software Specifications

The functionality have been validated on AWS p5.48xlarge EKS cluster (8x H100 80G).

Containers

We provide a multi-architecture container that is regularly updated. Use these containers to avoid dependency and environment issues.

• Latest container: ghcr.io/nvidia/jax:axlearn
• Nightly dated container: ghcr.io/nvidia/jax:axlearn-YYYY-MM-DD

When you start an interactive session:

• Navigate to /opt/axlearn inside the container.
• Place your persistent files in a mounted directory (e.g. /opt/axlearn/workspace).

Launching a container

Use the following command to launch a container:

$
docker run -ti --gpus=all --net=host --ipc=host -v <WORKSPACE_PATH>:/opt/axlearn/workspace -w /opt/axlearn <CONTAINER> /bin/bash

where WORKSPACE_PATH is the path to the directory where you would like to store any persistent files and container is the name of the AXLearn container. You can additionally add dataset and vocab paths with the -v flag.

Example 1: standard training for a Fuji model on EKS

AXLearn based models are called fuji. The fuji models come with several number of parameters: 1B, 3B, 7B and 70B. For each model there’s a V1, V2 and V3 version where:

• V1 exists for fuji-7B and fuji-70B. It specifies a vocab size of 32 * 1024, a max sequence length of 2048, and a total of 1T tokens for the 7B version and 1.4T for the 70B one;
• V2 exists for fuji-7B and fuji-70B. It specifies a vocab size of 32 * 1024, a max sequence length of 4096, and 2T tokens for both the 7B and 70B model;
• V3 is is used for 1B, 3B, 7B and 70B. The vocab size is 128*1024, the max sequence length 8192 and it provides 15T tokens for all the four models. You can check the above on the AXLearn code. Each model can then work in a different mode:
• -flash: uses flash attention;
• -flash-single-host: uses flash attention and it’s tuned to work on a single host. To run one of these models, on an EKS instance, you can follow this deployment file, with the running command depicted in the c4_trainer description, whose skeleton looks like:

1
apiVersion: batch/v1
2
kind: Job
3
metadata:
4
    name: axlearn-fuji
5
spec:
6
    completions: 1
7
    parallelism: 1
8
    template:
9
        spec:
10
            restartPolicy: Never
11
            containers:
12
                - name: axlearn-fuji-model
13
                  image: ghcr.io/nvidia/jax:axlearn
14
                  command:
15
                    - bash
16
                    - -xo
17
                    - pipefail
18
                    - -c
19
                    - |
20
                      BASEDIR="/opt/axlearn"
21
                      CONFIG="fuji-3B-v3-flash-single-host"
22
23
                      LOG_DIR=${BASEDIR}/logs
24
                      TRAINER_DIR=${LOG_DIR}/${CONFIG}-eks/trainer-dir
25
                      mkdir -p ${TRAINER_DIR}
26
27
                      python3 -m axlearn.common.launch_trainer_main \
28
                          --module=text.gpt.c4_trainer \
29
                          --config=${CONFIG} \
30
                          --trainer_dir=${TRAINER_DIR} \
31
                          --data_dir=gs://axlearn-public/tensorflow_datasets \
32
                          --jax_backend=gpu

This will run the fuji-3B-v3-flash-single-host model, and all the input configurations (e.g. max number of steps, sequence length, parallelism) can be found here. The input dataset is the public tensorflow C4 dataset.

Example 2: custom configuration training for a Fuji model on EKS

For specifying a custom configuration definition, we are using a Python script. The script is made based the following AXLearn c4 trainer script. The core configuration part is the following:

1
# Build the model config
2
config_fn = c4_trainer.named_trainer_configs()[config_name]
3
trainer_config: SpmdTrainer.Config = config_for_function(config_fn).fn()
4
# Intra-node parallelism
5
ici_mesh_shape = mesh_shape_from_axes(
6
        pipeline=ici_pp_size, data=ici_dp_size, fsdp=ici_fsdp_size, seq=ici_sqp_size
7
)
8
# Inter-node parallelism
9
dcn_mesh_shape = mesh_shape_from_axes(
10
        pipeline=dcn_pp_size, data=dcn_dp_size, fsdp=dcn_fsdp_size, seq=dcn_sqp_size
11
)
12
# Create a mesh
13
mesh_shape = HybridMeshShape(
14
        ici_mesh_shape=ici_mesh_shape, dcn_mesh_shape=dcn_mesh_shape
15
)
16
# GA & FSDP setup
17
mesh_rule = (
18
"custom",
19
ChainConfigModifier.default_config().set(
20
        config_modifiers=[
21
        GradientAccumulationModifier.default_config().set(
22
                grad_acc_steps=ga_size
23
        ),
24
        MeshShapeModifier.default_config().set(mesh_shape=mesh_shape),
25
        ]
26
),
27
)
28
trainer_config.mesh_rules = mesh_rule
29
trainer_config.mesh_shape = mesh_shape
30
# Max step
31
trainer_config.max_step = max_step
32
# Checkpoint directory
33
trainer_config.dir = trainer_dir
34
trainer_config.input.input_dispatcher.global_logical_batch_size = gbs_size
35
trainer_config.input.source.max_sequence_length = seq_len
36
trainer_config.checkpointer.save_policy.n = save_checkpoint_steps
37
trainer_config.checkpointer.keep_every_n_steps = save_checkpoint_steps
38
trainer_config.summary_writer.write_every_n_steps = write_summary_steps

After parsing the input parameters, config_fn = c4_trainer.named_trainer_configs()[config_name] retrieves the standard configuration for the model specified in config_name. The parallelism is define intra-node and inter-node. This function in AXLearn defines the construction of the mesh. Remember that for the intra-node mesh (ici_mesh), the parallelism nubmers product muber be as the same as the number of devices on a single node, while, for the inter-node mesh (dcn_mesh), the product must be equal to the number of nodes. The rest of the code specifies the gradient accumulation size, the global batch size, the max sequence length, and when to save checkpoints and summary files.

Then, we’re ready to launch the job with the following lines:

1
launch.setup()
2
trainer_config.set(
3
        recorder=config_for_function(lambda: measurement.global_recorder)
4
)
5
# Launch training
6
launch_trainer.run_trainer(
7
trainer_config=trainer_config,
8
)

In particular, launch.setup() refers to this code in AXLearn, where the main call to jax and its distributed initialization happens.

Testing

Container validation runs pytest against all tests in /opt/axlearn/axlearn/common using the test_axlearn.sh script bundled in the JAX-Toolbox container. See the repository’s CI workflows for example job definitions.