JAX Pathways

Course · Distributed trainingPhase 09 · Study guide

Is the distributed update still optimizing the same objective?

Systems Specialization · CPU ReferenceCPU reference; target qualification separate4 lessons (~6.3h)

Revision v0.7.0 ·

Start with the global mathematical operation, then specify ownership and communication. Logical CPU devices help expose placement and reduction mistakes before you spend accelerator time.

Preparation: Data & checkpoint recovery · Performance diagnosis

Runtime: Logical CPU devices for sharding practice; multi-host work needs a real cluster

Start the phaseSee the project

Check your starting point

Two devices see one and three valid examples. Should their mean gradients receive equal weight?

Compare your reasoning

No. Weight by valid example counts. Sum local gradient numerators and valid counts, then divide once. Padding and empty partitions must not silently change the denominator.

If this is unfamiliar, work through A sharded training step before attempting the project.

Build in stages

1. Verify placement and the global update

Inspect actual shard indices and compare each relevant output with an independent unpartitioned calculation, including padded and unequal partitions.

2. Connect communication and recovery

Inspect collectives without inferring network speed from payload size. Restore data position and optimizer state before comparing subsequent global steps.

Try a changed condition

Make a prediction before opening the reasoning. Keep your answer and the evidence you would collect.

A final batch leaves one partition empty. Should it contribute a zero mean, be omitted, or fail?

Compare an approach

Do not compute an undefined local mean. Accumulate a zero numerator and zero valid count for the empty partition, reduce globally and reject a globally empty batch. Compare the result with the valid unpartitioned examples.

Diagnose a likely failure

Symptom: The gradient grows with device count.

Check next: Check whether local means were summed without count normalization. Verify one complete update before interpreting a training curve.

Decide what is ready

Use sharded-training. Retain actual placements and collectives, independent uneven/empty-partition checks and full-state replay. Label logical CPU devices explicitly.

Where to put your project code

Write your code in: projects/sharded-training/my_model.py (copy projects/sharded-training/starter/model.py first: cp projects/sharded-training/starter/model.py projects/sharded-training/my_model.py).

Verify your implementation: python3 projects/sharded-training/tests/check.py --stage all --impl projects/sharded-training/my_model.py

Run the staged project checks on your implementation in projects/sharded-training/my_model.py and record your transfer experiment results in your portfolio.

Open A resilient sharded training experiment

What still needs further work

Real network collectives, multi-controller ingestion and host-loss recovery are unqualified. The optional launch instructions are preparation, not execution receipts.