Getting Started with AI task for ISC26 SCC (Virtual Part)

1. Introduction

About MaxText

MaxText is a high-performance, scalable, open-source Large Language Model (LLM) training framework developed by Google. Built on JAX, it is designed to be hardware-agnostic, running efficiently on both Google TPUs and NVIDIA GPUs. MaxText abstracts away much of the complexity of distributed training, allowing users to switch between parallelism strategies (Data Parallelism, FSDP, Tensor Parallelism, Pipeline Parallelism) with simple configuration flags. It serves as the reference implementation for Google's internal "AI Hypercomputer" stack.

About The Model: Qwen 3 14B

For this task, we will focus on Qwen 3 14B, a state-of-the-art dense model from Alibaba Cloud. MaxText supports the Qwen3 family natively.

• Parameters: ~14 Billion

• Architecture: Llama-like (RoPE, SwiGLU, RMSNorm, QK-Norm, GQA). See MaxText supported models.

• Why this model? At 14 billion parameters, it occupies a "middle ground" that benefits from multi-GPU setups and responds well to optimization (parallelism, batch size, rematerialization, etc.). Mastering this scale is critical for modern AI engineers.

Presentation:

Recording:

2. Environment Setup

2.1 Allocating Resources

Allocate a single node on the Romeo supercomputer. Each node is equipped with 4x NVIDIA GH200 (120GB) Grace-Hopper Superchips.

2.2 Launching the Container

We recommend using a pre-built image from NVIDIA. You can launch a container as follows:

The two -v flags mount host directories into the container:

• -v /scratch_p/$(whoami):/workspace — your personal scratch space becomes /workspace inside the container. This is where you will store outputs and run training.

• -v /project/r250119/maxtext:/share — the shared project directory (checkpoint, tokenizer, and dataset) is mounted at /share.

Verify that JAX sees your GPUs and maxtext is installed:

Note that no warning about cuDNN nor jaxlib should be present.

3. Preparing the Model

Before running any tests, you must download the Qwen weights and convert them to MaxText format.

3.1 Download Pre-converted Checkpoint and Tokenizer

The Qwen 3 14B weights have already been converted to MaxText's Orbax format and are available on shared storage.

• Checkpoint: /project/r250119/maxtext/model/qwen3_14b_orbax/0/items -- the converted Orbax weights (~26 GB).

• Tokenizer: /project/r250119/maxtext/model/qwen3_14b_hf -- the HuggingFace tokenizer directory (contains tokenizer.json). When passing tokenizer_path to MaxText, use the directory path, not the file.

3.2 Convert from HuggingFace Yourself (Optional)

If you want to run the conversion yourself to learn, follow these steps.

Warning: Installing PyTorch with CUDA support (torch+cu*) can corrupt the JAX CUDA runtime in this container. If you install it for conversion, restart the container afterwards before running any training.

4. Sanity Check: A Quick Test Run

Now that you have the environment and the model, verify everything works by running a 1-layer Qwen 3 14B on synthetic data. This checks that weights load and GPUs communicate without a full training run.

4.1 Run the Test

From the MaxText repo root (the maxtext directory after cloning), run the following. Include num_decoder_layers=1 so the override takes effect (the full model has many layers). This test run initialize the weight instead of loading from checkpointing and use synthetic dataset.

• Success: If you see a loss value printing and no crashes, your environment is ready.

• Next Step: For Task 1, use the full model (do not pass num_decoder_layers=1) and point your config or command-line args to the real dataset (see §5).

5. Competition Tasks

Objective: You are acting as an AI Engineer optimizing a training cluster. Your goal is to maximize Training Efficiency for the Qwen 3 14B model (high throughput in tokens/sec and model convergence). You may use any optimizations MaxText supports—e.g. parallelism, batch size, rematerialization, or other options—not only parallelism.

Task 1: Training Optimization

Continue pre-training the Qwen 3 14B model on your 4 GPUs using a real-world dataset. Use whatever optimizations you want (e.g. parallelism, batch size, rematerialization, or other MaxText options) to achieve the best throughput.

Dataset: We use a ~6-million-token subset of FineWeb-Edu (sample-10BT) from Hugging Face — a curated, high-quality web corpus filtered for educational content. The dataset is pre-downloaded as local Parquet files on shared storage.

• Location: /workspace/dataset/fineweb-edu/*.parquet

• Size: ~6M raw tokens

Constraints (fixed parameters):

• Training duration: Exactly 5 epochs (num_epoch=5). Set steps to a large value (e.g. steps=100000) — MaxText will stop automatically when the 5 epochs are complete.

• Sequence length: max_target_length=2048 (do not change this -- it defines the workload).

• Checkpoint & dataset: Use the provided Orbax checkpoint and FineWeb-Edu dataset.

• Stable loss: Final loss must be below 1.9.

Everything else is fair game: parallelism strategy, batch size, rematerialization policy, and any other MaxText options.

Running the Training: You can use a YAML config file or command-line overrides as in section 4. Here is a baseline command to get you started:

Note that this with this base configs, the run will take around 2 hours on a single node on Romeo.

Shell quoting note: The 'hf_train_files="*.parquet"' argument requires single quotes wrapping the entire key=value pair. Without them, OmegaConf interprets * as a YAML alias and the run will fail. If you use a YAML config file instead, write it as hf_train_files: "*.parquet" (no special quoting needed).

This runs the full 40-layer Qwen 3 14B with FSDP across 4 GPUs. Training will stop automatically after 5 passes over the dataset.

• Standard Output: Capture your logs! MaxText prints throughput (tokens/sec) and loss to stdout. Save to a file for your report: python3 ... | tee run.log.

• Checkpoints: With enable_checkpointing=true, MaxText saves checkpoints to <base_output_directory>/<run_name>/checkpoints/ (e.g. /workspace/output/task1_run/checkpoints/). You can use these to resume training if interrupted by passing load_parameters_path to the checkpoint directory.

• Pass Condition: Your training must be stable. Your final loss must be lower than 1.90 and showing a downward trend.

  • Note: With the pre-trained Qwen 3 14B weights on FineWeb-Edu, loss typically starts around ~2.3 and gradually decreases over the 5 epochs. The 2.50 threshold is a sanity check to ensure your run hasn't diverged.

  • Note: If your model diverges (loss spikes to NaN or increases), you receive 0 points for throughput.

Deliverables:

• Plot the Training Loss curve for your best run.

• Report your Average Throughput (tokens/device) and briefly explain the optimizations you used to get your best performance.

Task 2: Profiling & Analysis

Pick your best-performing configuration and profile a training step using the JAX xProf tool.

Step 1: Collect a profile trace. Re-run your best configuration with these additional flags (or add them to your YAML config). We run 10 steps and profile only the last one (step 10), which is a steady-state step after XLA compilation:

Replace ... with the same flags from your best Task 1 run. After the run completes, the profile trace will be saved under /workspace/myprofile/.

Step 2: Install and launch xProf. Inside the container:

This starts a web UI. If you are on a remote node, forward the port through SSH:

Then open http://localhost:8791/ in your browser. Select the run from the "Runs" dropdown, then choose trace_viewer under "Tools" to see the GPU timeline.

Alternatively, you can use TensorBoard with the profile plugin:

Step 3: Analyze the timeline and answer the following:

1. Communication Mapping: Can you identify the specific NCCL communication blocks in both the Forward and Backward passes? Which specific collective operation (e.g., All-Gather, Reduce-Scatter, All-Reduce) corresponds to which parallelism strategy (FSDP vs. DP) in your timeline?

2. Overlap Analysis: In an ideal high-performance setup, communication should happen concurrently with computation. Do you observe this overlap in your trace? If not, where are the stalls?

3. The "Funky" Factor: Look closely at the Backward Pass. Do you observe anything unusual or "funky" regarding the scheduling or execution order? (Hint: Look at how gradient computation interacts with weight updates or rematerialization). Share any other interesting bottlenecks or behaviors you discovered.

6. Bonus Task: The Multi-Node Challenge (Optional)

Scale beyond a single node: run the same Qwen 3 14B training across multiple nodes on Romeo.

The Challenge:

• Goal: Successfully complete a multi-node training run (e.g. 2 nodes, 8 GPUs total) and show that multi-node throughput is higher than your single-node run (demonstrating beneficial scaling).

• Focus: Demonstrate that your setup and config scale to more GPUs (communication, parallelism strategy, and job launch).

7. Submission & Grading

Submission Format

Each team will be given a submission directory on shared storage: /gpfs/projet/r250119/Submit/$USER/. Place the following inside:

1. Report (Markdown, Max 500 words + figures):

   • Loss and throughput plot(s) for Task 1.

   • Screenshots/Analysis of the xProf profile for Task 2.

   • Short explanation of your best run and the optimizations you used.

2. Log File:

   • run.log (or similar): raw stdout from your best 5-epoch run.

3. Config:

   • The config file or command-line setup used for your best run.

4. Final Checkpoint:

   • Copy the checkpoint directory from your best run (e.g. /workspace/output/task1_run/checkpoints/) into your submission folder. We will use this to verify your training results.

Evaluation Criteria (Total: 100 Points + 10 Bonus)

We evaluate your submission based on raw performance and analysis.

8. References

• MaxText configs: base.yml (in your clone, paths may use lowercase maxtext, e.g. src/maxtext/configs/base.yml).

• Qwen 3 14B: https://huggingface.co/Qwen/Qwen3-14B

9. Troubleshooting

Error: no space left on device when loading the container image (e.g. /tmp/container_images_oci.../....tar: no space left on device) Cause: The default TMPDIR (/tmp) does not have enough space to extract the container image. Fix: Point TMPDIR to a directory with more space before loading the image:

Error: httpx.LocalProtocolError: Illegal header value b'Bearer ' when using a Hugging Face dataset Cause: The training process is using an empty token (e.g. wrong or unset env var). MaxText/Hub expect HF_TOKEN. Fix: Pass the token in the same command so the Python process sees it: HF_TOKEN="$(cat ~/.cache/huggingface/token)" at the start of the line, then the rest of the train command. Or paste your token: HF_TOKEN="hf_xxxxxxxx" (from huggingface.co/settings/tokens). If it still fails, use Option A (synthetic) for the sanity check. For Task 1, the local parquet dataset does not require HF authentication.

Error: crun: create directory '/run/crun': No such file or directory (OCI runtime) Workaround:

Error: omegaconf.errors.GrammarParseError or Could not resolve interpolation when passing hf_train_files="*.parquet" Cause: OmegaConf interprets the * as a YAML alias/interpolation character. Fix: Wrap the entire argument in single quotes on the command line: 'hf_train_files="*.parquet"'. Alternatively, put this setting in a YAML config file where no shell quoting is needed:

Home directory is too small for storing docker/podman files?