AXLearn

AXLearn is a JAX-based deep learning library with a config system, trainer runtime, vision/audio experiment catalogs, and a GCP launcher/ops CLI.

Use this skill when the user asks about:

• axlearn.common configs, modules, trainers, inputs, learners, checkpointers, or launchers.
• axlearn CLI commands such as gcp config, gcp bundle, gcp launch, gcp vm, gcp bastion, gcp dataflow, gcp logs, or gcp auth.
• Vision workflows such as ImageNet, ResNet, CLIP, or other image-classification configs.
• Audio/ASR workflows such as Conformer, LibriSpeech, feature extraction, or WER evaluation.
• GPT / language-model trainer catalogs, tokenizers, MoE configs, or flash-attention paths.

Start here

1. Read references/repo-provenance.md if you need to check whether this skill is current for the checkout.
2. Read references/troubleshooting.md when installation, import, or optional dependency checks fail.
3. Use scripts/check_install.py for a safe import/version smoke check.
4. Route to the matching sub-skill:
   • sub-skills/training-core/ for trainer configs, fake-data smoke checks, and tokenizer setup.
   • sub-skills/language-models/ for GPT, MoE, flash-attention, and tokenizer catalog workflows.
   • sub-skills/cli-cloud/ for GCP config, bundle, launch, VM, bastion, Dataflow, logs, and auth.
   • sub-skills/vision-workflows/ for ResNet/ImageNet and other vision model recipes.
   • sub-skills/audio-asr/ for Conformer, LibriSpeech, and ASR workflows.

Installation and smoke check

For local inspection, install the editable package with the extras that match the workflow. Start with the base package, then add only the extras you need:

python -m pip install -e .
python -m pip install -e .[core,dev]

Common add-ons:

• audio for ASR workflows.
• gcp for cloud CLI workflows.
• orbax when checkpoint utilities are needed.
• dev only if you plan to run repo tests.

Minimal smoke checks:

python -I -c "import axlearn; print(axlearn.__file__)"
axlearn --help

If you are using the cloud CLI routes, also check:

axlearn gcp --help

Routing guidance

• Use training-core for local config construction, SpmdTrainer, launch_trainer_main, fake inputs, and short tutorial-style probes.
• Use language-models when the task names Fuji, Gala, Honeycrisp, Qwen, C4, Pajama, MoE, or flash attention.
• Use cli-cloud when the task names GCP activation, bundling, launching, bastion, Dataflow, logs, or auth.
• Use vision-workflows when the task names ImageNet, ResNet, image classification, or CLIP-like vision recipes.
• Use audio-asr when the task names LibriSpeech, Conformer, speech features, ASR, or WER.

If the task spans trainer config mechanics plus a domain family, start in training-core and then jump to the domain sub-skill.