← BACK TO LESSONPARALLEL HORIZONSREAD // 06

LESSON READING // 06

TRANSFORMER SCALING

HISTORICAL ANCHOR2022 · HOPPER FP8 / TRANSFORMER ENGINEREAD TIME · ABOUT 60 MIN

Combine low-precision numbers, intra-layer sharding, and communication overlap in one Transformer training step—and design numerical range alongside the parallel axis.

Write down answers as you go; the estimate includes code tracing and practice.

PLAY THE INTERACTIVE LESSON
COURSE PACINGABOUT 60 MINUTES
  1. 01ORIENT + OBJECTIVES05 MIN
  2. 02MENTAL MODEL08 MIN
  3. 03DEEP DIVES15 MIN
  4. 04CODE TRACE12 MIN
  5. 05HISTORY + ECOSYSTEM08 MIN
  6. 06PRACTICE + REVIEW12 MIN
ACCELERATED COMPUTING CAPABILITY MAPINDEPENDENT LEARNING COMPANION · NOT AN NVIDIA COURSE OR CERTIFICATION
PATH // 06
GPU WORKFLOWSPROFILINGMODEL PARALLELISM BRIDGE
THE OFFICIAL PATH EMPHASIZES

Compose CUDA Python, industry frameworks, and profiling into a reproducible workflow rather than optimizing an isolated kernel.

HOW THIS LESSON CONNECTS

The game combines FP8 recipes, tensor parallelism, and overlap; the lab validates numerical error, communication dependencies, and the timeline together.

SUGGESTED PREREQUISITES
  • Lessons 03–04
  • PyTorch basics
  • Transformer linear layers and backward pass
PRACTICE TOOLS
  • PyTorch
  • Transformer Engine
  • NCCL
  • Nsight Systems / Compute
01 · 05 MIN

ORIENT + OBJECTIVES

REMEMBER THIS FIRST

Fit numbers into range, shard the right dimension, and overlap only genuinely independent work.

BY THE END, YOU SHOULD BE ABLE TO

  1. 01

    Explain the roles of FP8 format, scale, amax, and high-precision master state.

  2. 02

    Choose parallel axes for wide layers, deep models, long sequences, and large batches.

  3. 03

    Draw compute/collective dependencies in tensor parallelism.

  4. 04

    Decide whether an asynchronous collective has a real, safe overlap window.

02 · 08 MIN

BUILD THE MENTAL MODEL

  1. 01
    SCALE

    Statistics such as amax produce a scale that maps a tensor into FP8's finite representable range.

  2. 02
    SHARD

    Tensor parallelism splits one layer; data, pipeline, and context parallelism split batch, depth, and sequence.

  3. 03
    OVERLAP

    After a collective launches, only work independent of its result is safe; consumers must wait.

03 · 15 MIN

CONCEPT DEEP DIVES

01

E4M3 AND E5M2 HAVE DIFFERENT JOBS

E4M3 retains more mantissa precision; E5M2 trades mantissa for wider exponent range. Hybrid recipes often separate forward and backward formats.

02

PARALLEL AXES SOLVE DIFFERENT LIMITS

TP targets wide layers, PP depth, CP long-sequence activations, and DP batch. The wrong axis leaves the original bottleneck intact.

03

MORE TP ALSO MEANS MORE COMMUNICATION

As local GEMMs shrink, the compute window available to hide communication can shrink too.

PART 01

LOW PRECISION STARTS WITH RANGE MANAGEMENT

E4M3 and E5M2 have limited range and precision. A scale maps source values into that range, often using amax. Too-small scales saturate large values; too-large scales collapse many small values to identical quantized values or zero.

Training distributions drift. Delayed and current scaling recipes update scale differently. Hybrid recipes commonly use E4M3 for forward weights/activations and wider-range E5M2 for gradients. Not every operation belongs in FP8, and high-precision state can remain necessary.

PAUSE AND REASONWhy is one fixed scale for the whole training run risky?

REFERENCE ANSWERTensor distributions drift, so the fixed scale can gradually cause saturation or poor small-value resolution.

PART 02

PARALLEL AXES SOLVE ORTHOGONAL PROBLEMS

Data parallelism splits batch and replicates the model; tensor parallelism splits layer tensors; pipeline parallelism splits depth; context parallelism splits sequence and related activations. They compose, but every axis adds communication, scheduling, or state management.

Choose from the bottleneck: TP for a wide layer, PP for depth, CP for long-sequence activations, and DP for throughput. Rank groups must keep each collective inside the correct communication domain.

PAUSE AND REASONDoes more DP directly fix activation OOM from a long context?

REFERENCE ANSWERUsually no. Each DP rank still handles its full sequence; investigate context parallelism, sequence sharding, or checkpointing.

PART 03

ASYNC IS A HANDLE; INDEPENDENT WORK IS THE WINDOW

An async All-Gather returns a handle, but an immediate wait leaves communication fully exposed. Overlap requires computation that does not consume the gathered result and can run concurrently without destructive resource contention.

Higher TP shrinks each rank's GEMM while increasing communication share, potentially shortening the overlap window. Confirm the timeline, waits, streams, message size, and topology in a profiler.

PAUSE AND REASONHow different is async All-Gather followed immediately by wait from synchronous All-Gather?

REFERENCE ANSWERUsually very little; no independent work occupies the interval, so the API did not create overlap.

04 · 12 MIN

PUT IT BACK INTO CODE

PROGRAM MODELManage FP8 with a recipe, then wait at the dependency boundary
01recipe = DelayedScaling(fp8_format=Format.HYBRID)02with te.autocast(enabled=True, recipe=recipe):03  local = column_parallel_linear(x)04  handle = all_gather_async(local)05  independent = other_linear(x)06  gathered = handle.wait()

TRACE THE CODE

  1. 01
    RECIPEDECLARE FORMAT + SCALING POLICY

    The recipe centralizes statistics and scale updates, while model-level accuracy still needs validation.

  2. 02
    LOCALEACH TP RANK COMPUTES A LOCAL LINEAR

    Row/column sharding defines local ownership and the collective that follows.

  3. 03
    ASYNCLAUNCH COMMUNICATION AND KEEP A HANDLE

    The result is not ready to consume; buffers and streams must remain valid.

  4. 04
    WAITWAIT AT THE FIRST CONSUMER

    Place truly independent computation in the window, then wait before reading gathered data.

REF · 06

KNOWLEDGE ATLAS

OPTIONAL REFERENCE · REVISIT AS NEEDED · OUTSIDE THE CORE 60 MINUTES

CORE VOCABULARY

FP8
A family of 8-bit floating formats whose limited range requires explicit scaling and higher-precision retention or accumulation paths.
amax
The maximum absolute value in a tensor or history window, used to estimate scale and clipping risk.
Delayed Scaling
Uses historical amax to estimate the current scale and avoid an extra pre-quantization read, trading in scale lag.
Tensor Parallel
Splits linear/attention tensors across ranks, then combines local results with collectives.
Overlap
Advancing communication alongside truly independent computation; an async call alone does not guarantee it.
EXAMPLE

WORKED EXAMPLE

  1. 01
    COMPUTE SCALE

    scale ≈ 448 ÷ 240 ≈ 1.87; an original value of 120 maps to about 224.

  2. 02
    WATCH A SPIKE

    A sudden value of 300 maps to about 560 and clips at 448.

  3. 03
    COMPUTE STEP TIME

    T_step ≈ 12 + 5 − 3 = 14 ms; uncovered communication remains critical.

RESULT

One step can both clip an outlier because of stale scale and pay about 2 ms for communication not hidden by compute.

WHY IT MATTERS

Low-precision training needs numerical evidence and a distributed timeline—not throughput or loss alone.

DIAGNOSTIC PLAYBOOK

01SYMPTOM

FP8 loss spikes or becomes NaN

INSPECT FIRST
amax/scale history, clipping, sensitive layers, master weights, and reduction precision
EVIDENCE
Per-layer amax, overflow counts, BF16 control, and a multi-step loss curve
02SYMPTOM

A larger TP degree slows the step

INSPECT FIRST
Tiny local GEMMs, collective share, topology, and waits
EVIDENCE
Rank-level compute/communication spans across a TP sweep
03SYMPTOM

Async collective is not hidden

INSPECT FIRST
True downstream independence, shared resources, and early waits
EVIDENCE
Collective, GEMM, and wait overlap in Nsight Systems

HARDWARE → ECOSYSTEM

  1. 01FP32 TRAINING

    Uniform high precision simplified numerical reasoning.

    Memory, bandwidth, and matrix throughput were expensive.
  2. 02FP16/BF16 MIXED PRECISION

    Low-precision Tensor paths paired with high-precision accumulation and loss scaling.

    Numerical policy became performance engineering.
  3. 03HOPPER FP8 + TRANSFORMER ENGINE

    Hardware formats, scale recipes, and framework modules coordinated.

    Optimization expanded to layer statistics, process groups, and communication scheduling.
05 · 08 MIN

HARDWARE + ECOSYSTEM COORDINATE

2022 · HOPPER TRANSFORMER ENGINE

Hopper Tensor Cores added FP8 matrix paths. Transformer Engine manages FP8-safe operations, scales, and amax history in software. Frameworks such as Megatron Core compose TP, PP, DP, CP, and overlap into configurable training systems.

06 · 12 MIN

PRACTICE + REVIEW

A single extremely wide Linear layer does not fit one GPU. Which axis should you investigate first?

  1. A. Tensor parallelism
  2. B. Only more data parallelism
  3. C. Only batch size 1
REVEAL ANSWER
A

Tensor parallelism directly shards intra-layer weights and compute. Standard data parallelism still replicates the full layer on every rank.

02FORMAT

Why are gradients often associated with E5M2 and forward tensors with E4M3?

HINT

Compare exponent range and mantissa precision.

REFERENCE ANSWER

Gradients can need wider range, while forward weights/activations often value extra precision. E5M2 trades mantissa for range; E4M3 does the reverse.

03SELECT

The model is very deep, each layer fits, but the full model does not. Which axis first?

HINT

Split by layer depth.

REFERENCE ANSWER

Pipeline parallelism, followed by analysis of bubbles, microbatches, and activation transfer.

04DEPENDENCY

Work after async All-Gather reads part of its output. Can it overlap safely?

HINT

Is the work genuinely independent?

REFERENCE ANSWER

Not directly. It must wait for required data unless the algorithm supports finer-grained completion and consumption.

05DIAGNOSE

TP 8 is slower than TP 4. What should you compare?

HINT

Local compute shrinks while communication grows.

REFERENCE ANSWER

Compare per-rank GEMM shape/efficiency, collective size/time, overlap window, topology, and wait ratio.

LAB · 45–70 MIN

OPTIONAL HANDS-ON LAB

OUTSIDE THE CORE 60 MINUTES · REQUIRES A SUITABLE CUDA / GPU ENVIRONMENT

LAB GOAL

VALIDATE ONE FP8 + TENSOR-PARALLEL TRAINING STEP

Compare BF16 with an FP8 recipe and inspect whether TP communication overlaps genuinely independent work.

PROCEDURE

  1. 01

    Build a seeded BF16 reference with loss and gradient summaries.

  2. 02

    Enable a Transformer Engine recipe and record format, scaling policy, and hardware support.

  3. 03

    Compare TP=1 and TP>1 with rank shapes, collectives, and error.

  4. 04

    Profile async collectives, independent linear work, and waits.

EVIDENCE OF COMPLETION

Submit environment capability, BF16/FP8 numerics, TP rank mapping, and a timeline; on unsupported hardware, provide a labeled static recipe/shape trace.

STRETCH CHALLENGE

Sweep TP degree and graph local GEMM, collective time, and overlap to find over-sharding.

OFFICIAL LEARNING ENTRY POINTSTransformer Engine DocumentationMegatron Core Parallelism Guide
REF

GO DEEPER

01
Transformer Engine · FP8 Current ScalingE4M3, E5M2, scales, and hybrid recipes
OFFICIAL SOURCE · OPENS IN A NEW TAB02
Megatron Core · Parallelism StrategiesObjectives and composition of DP, TP, PP, and CP
OFFICIAL SOURCE · OPENS IN A NEW TAB