「技术」LLM One Page: nanochat speedrun

1
#!/bin/bash
2

3
# This script is configured to train your own GPT-2 grade LLM (pretraining + finetuning)
4
# It is designed to run on a blank 8XA800 GPU node and takes approximately 3 hours to complete.
5

6
# 1) Example launch (simplest):
7
# bash runs/speedrun.sh
8
# 2) Example launch in a screen session (because the run takes ~3 hours):
9
# screen -L -Logfile runs/speedrun.log -S speedrun bash runs/speedrun.sh
10
# 3) Example launch with wandb logging, but see below for setting up wandb first:
11
# WANDB_RUN=speedrun screen -L -Logfile runs/speedrun.log -S speedrun bash runs/speedrun.sh
12

13
# Default intermediate artifacts directory is in ~/.cache/nanochat
14
export OMP_NUM_THREADS=1
15
export NANOCHAT_BASE_DIR="$HOME/.cache/nanochat"
16
mkdir -p $NANOCHAT_BASE_DIR
17

18
# -----------------------------------------------------------------------------
19
# Python venv setup with uv
20

21
# install uv (if not already installed)
22
command -v uv &> /dev/null || curl -LsSf https://astral.sh/uv/install.sh | sh
23
# create a .venv local virtual environment (if it doesn't exist)
24
[ -d ".venv" ] || uv venv
25
# install the repo dependencies
26
uv sync --extra gpu
27
# activate venv so that `python` uses the project's venv instead of system python
28
source .venv/bin/activate
29

30
# -----------------------------------------------------------------------------
31
# wandb setup
32
# If you wish to use wandb for logging (it's nice!, recommended).
33
# 1) Make sure to first log in to wandb, e.g. run:
34
#    `wandb login`
35
# 2) Set the WANDB_RUN environment variable when running this script, e.g.:
36
#    `WANDB_RUN=d26 bash speedrun.sh`
37
if [ -z "$WANDB_RUN" ]; then
38
    # by default use "dummy" : it's handled as a special case, skips logging to wandb
39
    WANDB_RUN=dummy
40
fi
41

42
# -----------------------------------------------------------------------------
43
# During the course of the run, we will be writing markdown reports to the report/
44
# directory in the base dir. This command clears it out and writes a header section
45
# with a bunch of system info and a timestamp that marks the start of the run.
46
python -m nanochat.report reset
47

48
# -----------------------------------------------------------------------------
49
# Tokenizer
50

51
# Download the first ~2B characters of pretraining dataset
52
# each data shard is ~250M chars
53
# so we download 2e9 / 250e6 = 8 data shards at this point
54
# each shard is ~100MB of text (compressed), so this is about ~800MB of data on disk
55
# look at dev/repackage_data_reference.py for details on how this data was prepared
56
python -m nanochat.dataset -n 8
57
# Immediately also kick off downloading more shards in the background while tokenizer trains
58
# Approximately 350 shards are needed for 10B tokens of data for pretraining.
59
# The maximum total number of shards available in the entire dataset is 1822.
60
python -m nanochat.dataset -n 370 &
61
DATASET_DOWNLOAD_PID=$!
62
# train the tokenizer with vocab size 2**15 = 32768 on ~2B characters of data
63
python -m scripts.tok_train
64
# evaluate the tokenizer (report compression ratio etc.)
65
python -m scripts.tok_eval
66

67
# -----------------------------------------------------------------------------
68
# Base model (pretraining)
69
echo "Waiting for dataset download to complete..."
70
wait $DATASET_DOWNLOAD_PID
71

72
# d24 model (slightly overtrained is enough to beat GPT-2 => increase data:params ratio from compute optimal 10.5 (default) to 12)
73
# MODIFIED FOR A800: Removed --fp8 flag
74
torchrun --standalone --nproc_per_node=8 -m scripts.base_train -- --depth=26 --target-param-data-ratio=8.25 --device-batch-size=16 --run=$WANDB_RUN
75
# evaluate the model: CORE metric, BPB on train/val, and draw samples
76
torchrun --standalone --nproc_per_node=8 -m scripts.base_eval -- --device-batch-size=16
77

78
# -----------------------------------------------------------------------------
79
# SFT (teach the model conversation special tokens, tool use, multiple choice)
80

81
# download 2.3MB of synthetic identity conversations to impart a personality to nanochat
82
# see dev/gen_synthetic_data.py for details on how this data was prepared and to get a sense of how you can easily tune it
83
curl -L -o $NANOCHAT_BASE_DIR/identity_conversations.jsonl https://karpathy-public.s3.us-west-2.amazonaws.com/identity_conversations.jsonl
84

85
# run SFT and eval the model
86
torchrun --standalone --nproc_per_node=8 -m scripts.chat_sft -- --device-batch-size=16 --run=$WANDB_RUN
87
torchrun --standalone --nproc_per_node=8 -m scripts.chat_eval -- -i sft
88

89
# chat with the model over CLI! Leave out the -p to chat interactively
90
# python -m scripts.chat_cli -p "Why is the sky blue?"
91

92
# even better, chat with your model over a pretty WebUI ChatGPT style
93
# python -m scripts.chat_web
94

95
# -----------------------------------------------------------------------------
96
# Generate the full report by putting together all the sections
97
# report.md is the output and will be copied to current directory for convenience
98
python -m nanochat.report generate

1
        // Stream ingestion loop: refill under GIL, process without GIL (parallel)
2
      loop {
3
          let exhausted = refill(&mut buf)?;
4
          if buf.is_empty() && exhausted {
5
              break;
6
          }
7

8
          total_sequences += buf.len() as u64;
9

10
          let pattern = self.compiled_pattern.clone();
11
          let local: AHashMap<CompactString, i32> = py.detach(|| {
12
              buf.par_iter()
13
                  .map(|s| {
14
                      let mut m: AHashMap<CompactString, i32> = AHashMap::new();
15
                      for mat in pattern.find_iter(s) {
16
                          let piece = mat.expect("regex match failed").as_str();
17
                          *m.entry(CompactString::from(piece)).or_default() += 1;
18
                      }
19
                      m
20
                  })
21
                  .reduce(AHashMap::new, |mut a, b| {
22
                      for (k, v) in b {
23
                          *a.entry(k).or_default() += v;
24
                      }
25
                      a
26
                  })
27
          });
28

29
          // Merge local into global (single-threaded)
30
          for (k, v) in local {
31
              *counts.entry(k).or_default() += v;
32
          }
33

34
          if exhausted {
35
              break;
36
          }
37
      }

1
impl Tokenizer {
2
    /// Core incremental BPE training given unique words and their counts.
3
    /// `words`: one entry per unique chunk (Vec<u32> of token-ids/bytes).
4
    /// `counts`: same length as `words`, count per chunk.
5
    fn train_core_incremental(&mut self, mut words: Vec<Word>, counts: Vec<i32>, vocab_size: u32) {
6
        assert!(vocab_size >= 256, "vocab_size must be at least 256");
7
        let num_merges = vocab_size - 256;
8
        log::info!("Starting BPE training: {} merges to compute", num_merges);
9
        self.merges.clear();
10

11
        // ---- Initial pair_counts and where_to_update (parallel) ----
12
        log::info!(
13
            "Computing initial pair counts from {} unique sequences",
14
            words.len()
15
        );
16
        let (mut pair_counts, mut where_to_update) = count_pairs_parallel(&words, &counts);
17

18
        // ---- Build heap ----
19
        log::info!("Building heap with {} unique pairs", pair_counts.len());
20
        let mut heap = OctonaryHeap::with_capacity(pair_counts.len());
21
        for (pair, pos) in where_to_update.drain() {
22
            let c = *pair_counts.get(&pair).unwrap_or(&0);
23
            if c > 0 {
24
                heap.push(MergeJob {
25
                    pair,
26
                    count: c as u64,
27
                    pos,
28
                });
29
            }
30
        }
31

32
        // ---- Merge loop ----
33
        log::info!("Starting merge loop");
34
        let mut merges_done = 0u32;
35
        let mut last_log_percent = 0u32;
36

37
        while merges_done < num_merges {
38
            let Some(mut top) = heap.pop() else {
39
                break;
40
            };
41

42
            // Lazy refresh: if the count changed since we queued this job, update and requeue
43
            let current = *pair_counts.get(&top.pair).unwrap_or(&0);
44
            if current <= 0 {
45
                // Pair no longer exists or has non-positive count, skip it
46
                continue;
47
            }
48
            if top.count != current as u64 {
49
                top.count = current as u64;
50
                heap.push(top);
51
                continue;
52
            }
53

54
            // Record merge
55
            let new_id = 256 + merges_done;
56
            self.merges.insert(top.pair, new_id);
57

58
            // Merge this pair in all words where it occurs
59
            let mut local_pos_updates: AHashMap<Pair, AHashSet<usize>> = AHashMap::new();
60
            for &word_idx in &top.pos {
61
                // Apply merge to this word and collect pair-count deltas
62
                let changes = words[word_idx].merge_pair(top.pair, new_id);
63
                // Update global pair counts based on this word's count
64
                for (pair, delta) in changes {
65
                    let delta_total = delta * counts[word_idx];
66
                    if delta_total != 0 {
67
                        *pair_counts.entry(pair).or_default() += delta_total;
68
                        if delta > 0 {
69
                            local_pos_updates.entry(pair).or_default().insert(word_idx);
70
                        }
71
                    }
72
                }
73
            }
74

75
            // Add the updated pair counts back to the heap
76
            for (pair, pos) in local_pos_updates {
77
                let cnt = *pair_counts.get(&pair).unwrap_or(&0);
78
                if cnt > 0 {
79
                    heap.push(MergeJob {
80
                        pair,
81
                        count: cnt as u64,
82
                        pos,
83
                    });
84
                }
85
            }
86

87
            merges_done += 1;
88

89
            // Log progress every 1%
90
            let current_percent = (merges_done * 100) / num_merges;
91
            if current_percent > last_log_percent {
92
                log::info!(
93
                    "Progress: {}% ({}/{} merges) - Last merge: {:?} -> {} (frequency: {})",
94
                    current_percent,
95
                    merges_done,
96
                    num_merges,
97
                    top.pair,
98
                    new_id,
99
                    top.count
100
                );
101
                last_log_percent = current_percent;
102
            }
103
        }
104

105
        log::info!("Finished training: {} merges completed", merges_done);
106
    }
107
}

1
    def forward(self, idx, targets=None, kv_cache=None, loss_reduction='mean'):
2
        B, T = idx.size()
3

4
        # Grab the rotary embeddings for the current sequence length (they are of shape (1, seq_len, 1, head_dim/2))
5
        assert T <= self.cos.size(1), f"Sequence length grew beyond the rotary embeddings cache: {T} > {self.cos.size(1)}"
6
        assert idx.device == self.cos.device, f"Rotary embeddings and idx are on different devices: {idx.device} != {self.cos.device}"
7
        assert self.cos.dtype == torch.bfloat16, "Rotary embeddings must be in bfloat16"
8
        # if kv cache exists, we need to offset the rotary embeddings to the current position in the cache
9
        T0 = 0 if kv_cache is None else kv_cache.get_pos()
10
        cos_sin = self.cos[:, T0:T0+T], self.sin[:, T0:T0+T] # truncate cache to current sequence length
11

12
        # Forward the trunk of the Transformer
13
        x = self.transformer.wte(idx) # embed current token
14
        x = norm(x)
15
        x0 = x  # save initial normalized embedding for x0 residual
16
        for i, block in enumerate(self.transformer.h):
17
            x = self.resid_lambdas[i] * x + self.x0_lambdas[i] * x0
18
            ve = self.value_embeds[str(i)](idx) if str(i) in self.value_embeds else None
19
            x = block(x, ve, cos_sin, self.window_sizes[i], kv_cache)
20
        x = norm(x)
21

22
        # Forward the lm_head (compute logits)
23
        softcap = 15 # smoothly cap the logits to the range [-softcap, softcap]
24
        logits = self.lm_head(x) # (B, T, padded_vocab_size) <- very big tensor, large amount of memory
25
        logits = logits[..., :self.config.vocab_size] # slice to remove padding
26
        logits = logits.float() # switch to fp32 for logit softcap and loss computation
27
        logits = softcap * torch.tanh(logits / softcap) # squash the logits
28

29
        if targets is not None:
30
            # training: given the targets, compute and return the loss
31
            # TODO experiment with chunked cross-entropy?
32
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1, reduction=loss_reduction)
33
            return loss
34
        else:
35
            # inference: just return the logits directly
36
            return logits

1
class CausalSelfAttention(nn.Module):
2
    def __init__(self, config, layer_idx):
3
        super().__init__()
4
        self.layer_idx = layer_idx
5
        self.n_head = config.n_head
6
        self.n_kv_head = config.n_kv_head
7
        self.n_embd = config.n_embd
8
        self.head_dim = self.n_embd // self.n_head
9
        assert self.n_embd % self.n_head == 0
10
        assert self.n_kv_head <= self.n_head and self.n_head % self.n_kv_head == 0
11
        self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
12
        self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
13
        self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
14
        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
15
        self.ve_gate_channels = 32
16
        self.ve_gate = nn.Linear(self.ve_gate_channels, self.n_kv_head, bias=False) if has_ve(layer_idx, config.n_layer) else None
17

18
    def forward(self, x, ve, cos_sin, window_size, kv_cache):
19
        B, T, C = x.size()
20

21
        # Project the input to get queries, keys, and values
22
        # Shape: (B, T, H, D) - FA3's native layout, no transpose needed!
23
        q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
24
        k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
25
        v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
26

27
        # Value residual (ResFormer): mix in value embedding with input-dependent gate per head
28
        if ve is not None:
29
            ve = ve.view(B, T, self.n_kv_head, self.head_dim)
30
            gate = 2 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))  # (B, T, n_kv_head), range (0, 2)
31
            v = v + gate.unsqueeze(-1) * ve
32

33
        # Apply Rotary Embeddings to queries and keys to get relative positional encoding
34
        cos, sin = cos_sin
35
        q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin)
36
        q, k = norm(q), norm(k) # QK norm
37

38
        # Flash Attention (FA3 on Hopper+, PyTorch SDPA fallback elsewhere)
39
        # window_size is (left, right) tuple: (N, 0) for causal, (-1, 0) for full context
40
        if kv_cache is None:
41
            # Training: causal attention with optional sliding window
42
            y = flash_attn.flash_attn_func(q, k, v, causal=True, window_size=window_size)
43
        else:
44
            # Inference: use flash_attn_with_kvcache which handles cache management
45
            k_cache, v_cache = kv_cache.get_layer_cache(self.layer_idx)
46
            y = flash_attn.flash_attn_with_kvcache(
47
                q, k_cache, v_cache,
48
                k=k, v=v,
49
                cache_seqlens=kv_cache.cache_seqlens,
50
                causal=True,
51
                window_size=window_size,
52
            )
53
            # Advance position after last layer processes
54
            if self.layer_idx == kv_cache.n_layers - 1:
55
                kv_cache.advance(T)
56

57
        # Re-assemble the heads and project back to residual stream
58
        y = y.contiguous().view(B, T, -1)
59
        y = self.c_proj(y)
60
        return y

1
class CausalSelfAttention(nn.Module):
2
    def __init__(self, config, layer_idx):
3
        super().__init__()
4
        self.layer_idx = layer_idx
5
        self.n_head = config.n_head           # Query 的头数 (h)
6
        self.n_kv_head = config.n_kv_head     # Key/Value 的头数 (g)
7
        self.n_embd = config.n_embd           # 嵌入维度 (C 或 d_model)
8

9
        # 每个头的维度：C = n_head * head_dim，确保多头拼接后能还原回 C
10
        self.head_dim = self.n_embd // self.n_head
11

12
        # 健壮性检查：确保嵌入维度能被头数整除
13
        assert self.n_embd % self.n_head == 0
14

15
        # GQA 核心约束：KV 头数必须小于等于 Q 头数，且能被整除（实现均匀分组）
16
        # 如果 n_kv_head = 1，则是 MQA；如果 n_kv_head = n_head，则是 MHA
17
        assert self.n_kv_head <= self.n_head and self.n_head % self.n_kv_head == 0
18

19
        # Q 投影层：输出维度为 n_head * head_dim，即保持完整的 C
20
        self.c_q = nn.Linear(self.n_embd, self.n_head * self.head_dim, bias=False)
21

22
        # K/V 投影层（GQA 的精髓）：输出维度大大缩小，仅为 n_kv_head * head_dim
23
        # 这里的输出宽度直接决定了 KV Cache 的显存占用量
24
        self.c_k = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
25
        self.c_v = nn.Linear(self.n_embd, self.n_kv_head * self.head_dim, bias=False)
26

27
        # 输出投影层 (Wo)：负责将所有头拼接后的结果融合，重新映射回 n_embd 空间
28
        self.c_proj = nn.Linear(self.n_embd, self.n_embd, bias=False)
29
# ...

1
    def forward(self, x, ve, cos_sin, window_size, kv_cache):
2
        B, T, C = x.size()
3

4
        # Project the input to get queries, keys, and values
5
        # Shape: (B, T, H, D) - FA3's native layout, no transpose needed!
6

7
        # q 被拆成 n_head 个头，每个头的维度是 head_dim
8
        # k,v 被分成 g 组，每个组有 n_kv_head 个头，每个头的维度是 head_dim， g = n_head // n_kv_head
9
        q = self.c_q(x).view(B, T, self.n_head, self.head_dim)
10
        k = self.c_k(x).view(B, T, self.n_kv_head, self.head_dim)
11
        v = self.c_v(x).view(B, T, self.n_kv_head, self.head_dim)
12

13
        # Value residual (ResFormer): mix in value embedding with input-dependent gate per head
14
        # ...
15

16
        # Apply Rotary Embeddings to queries and keys to get relative positional encoding
17
        # ...
18

19
        # Flash Attention (FA3 on Hopper+, PyTorch SDPA fallback elsewhere)
20
        # 训练 / 推理 分别调用 flash-attention 算子，其中推理有 kv cache
21

22
        # Re-assemble the heads and project back to residual stream
23

24
        # 将输出 y 的形状重新转换为 (B, T, C)
25
        y = y.contiguous().view(B, T, -1)
26

27
        # c_proj 是一个全连接层，它让来自不同头的特征进行第一次实质性的加权混合
28
        # 同时还将注意力层的输出通过线性变换，调整到适合与原始输入 $x$ 进行**残差相加（Residual Add）**的状态
29
        y = self.c_proj(y)
30
        return y

1
# flash_attn_func ...
2
return _fa3.flash_attn_func(q, k, v, causal=causal, window_size=window_size)
3

4
# sdpa_attention ...
5
return F.scaled_dot_product_attention(q, k, v, attn_mask=mask, enable_gqa=enable_gqa)

1
def apply_rotary_emb(x, cos, sin):
2
    assert x.ndim == 4  # multihead attention
3
    d = x.shape[3] // 2
4
    x1, x2 = x[..., :d], x[..., d:] # split up last dim into two halves
5
    y1 = x1 * cos + x2 * sin # rotate pairs of dims (正负号这里不重要，看起来是实现的顺时针旋转)
6
    y2 = x1 * (-sin) + x2 * cos
7
    return torch.cat([y1, y2], 3)
8

9
# ...
10

11
def _precompute_rotary_embeddings(self, seq_len, head_dim, base=10000, device=None):
12
    # TODO: bump base theta more? e.g. 100K is more common more recently
13
    # autodetect the device from model embeddings
14
    if device is None:
15
        device = self.transformer.wte.weight.device
16
    # stride the channels
17
    channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
18
    inv_freq = 1.0 / (base ** (channel_range / head_dim))
19
    # stride the time steps
20
    t = torch.arange(seq_len, dtype=torch.float32, device=device)
21
    # calculate the rotation frequencies at each (time, channel) pair
22
    freqs = torch.outer(t, inv_freq)
23
    cos, sin = freqs.cos(), freqs.sin()
24
    cos, sin = cos.bfloat16(), sin.bfloat16() # keep them in bfloat16
25
    cos, sin = cos[None, :, None, :], sin[None, :, None, :] # add batch and head dims for later broadcasting
26
    return cos, sin

1
def has_ve(layer_idx, n_layer):
2
    """Returns True if GPT layer should have Value Embedding (alternating, last layer always included)."""
3
    return layer_idx % 2 == (n_layer - 1) % 2
4

5
# ...
6

7
        self.ve_gate_channels = 32
8
        self.ve_gate = nn.Linear(self.ve_gate_channels, self.n_kv_head, bias=False) if has_ve(layer_idx, config.n_layer) else None
9

10
# ...
11

12
        # Value residual (ResFormer): mix in value embedding with input-dependent gate per head
13
        if ve is not None:
14
            ve = ve.view(B, T, self.n_kv_head, self.head_dim)
15
            gate = 2 * torch.sigmoid(self.ve_gate(x[..., :self.ve_gate_channels]))  # (B, T, n_kv_head), range (0, 2)
16
            v = v + gate.unsqueeze(-1) * ve

1
class MLP(nn.Module):
2
    def __init__(self, config):
3
        super().__init__()
4
        self.c_fc = nn.Linear(config.n_embd, 4 * config.n_embd, bias=False)
5
        self.c_proj = nn.Linear(4 * config.n_embd, config.n_embd, bias=False)
6

7
    def forward(self, x):
8
        x = self.c_fc(x)
9
        x = F.relu(x).square()
10
        x = self.c_proj(x)
11
        return x

1
        x = norm(x)
2

3
        # Forward the lm_head (compute logits)
4
        softcap = 15 # smoothly cap the logits to the range [-softcap, softcap]
5
        logits = self.lm_head(x) # (B, T, padded_vocab_size) <- very big tensor, large amount of memory
6
        logits = logits[..., :self.config.vocab_size] # slice to remove padding
7
        logits = logits.float() # switch to fp32 for logit softcap and loss computation
8
        logits = softcap * torch.tanh(logits / softcap) # squash the logits
9

10
        if targets is not None:
11
            # training: given the targets, compute and return the loss
12
            # TODO experiment with chunked cross-entropy?
13
            loss = F.cross_entropy(logits.view(-1, logits.size(-1)), targets.view(-1), ignore_index=-1, reduction=loss_reduction)
14
            return loss
15
        else:
16
            # inference: just return the logits directly
17
            return logits

1
while True:
2
# ...
3
            rg_idx = ddp_rank
4
        while rg_idx < pf.num_row_groups:
5
            rg = pf.read_row_group(rg_idx)
6
            batch = rg.column('text').to_pylist()
7
            for i in range(0, len(batch),tokenizer_batch_size):
8
                yield batch[i:i+tokenizer_batch_size],(pq_idx, rg_idx, epoch)
9
            rg_idx += ddp_world_size
10
        pq_idx += 1
11
    first_pass = False
12
    epoch += 1

1
# ...
2
            grad = p.grad
3
        if p.numel() < 1024:
4
            # Small params: all_reduce (no scattergather needed)
5
            future = dist.all_reduce(grad, op=distReduceOp.AVG, async_op=True).get_future()
6
            param_infos[p] = dict(future=future,grad_slice=grad, is_small=True)

1
else:
2
    # Large params: reduce_scatter
3
    assert grad.shape[0] % world_size == 0, f"AdamW reduce_scatter requires shape[0] ({grad.shape[0]}) divisible by world_size ({world_size})"
4
    rank_size = grad.shape[0] // world_size
5
    # rank_size 算出了“每张显卡应该负责更新多大的参数切片”。因为有 4 张卡，所以 4096 行的大矩阵被平分为 4 份，每张卡负责 1024 行（即 (1024, 768) 大小的切片）。
6
    grad_slice = torch.empty_like(grad[:rank_size])
7
    future = dist.reduce_scatter_tensor(grad_slice, grad, op=dist.ReduceOp.AVG, async_op=True).get_future()
8
    '''
9
    dist.reduce_scatter_tensor ：这是一个底层集合通信函数（通常由 NCCL 硬件加速）。
10
    - 它把所有 4 张卡的 4096 行的大梯度矩阵 grad 收上来， 先在底层网络中把 4 个人的大矩阵对应位置相加求平均（ op=dist.ReduceOp.AVG ） 。
11
    - 切分结果 ：求完平均后，一个完整的 4096 行平均梯度矩阵诞生了。但它 不把整个结果发给所有人 （那是 All-Reduce）！
12
    - 分发结果（Scatter） ：它把这个 4096 行的平均结果切成 4 份（每份 1024 行）：
13
        - 把第 1 份（0~1023行）发给 GPU 0。
14
        - 把第 2 份（1024~2047行）发给 GPU 1。
15
        - 把第 3 份发给 GPU 2，第 4 份发给 GPU 3。
16
    '''
17
    param_infos[p] = dict(future=future, grad_slice=grad_slice, is_small=False)

1
else:
2
    rank_size = p.shape[0] // world_size
3
    p_slice = p[rank * rank_size:(rank + 1) * rank_siz
4
# ...
5
# 执行 AdamW 更新步骤 (此时状态 state['exp_avg'] 也只是这个切片大小
6
adamw_step_fused(
7
    p_slice, grad_slice, state['exp_avg'], state['exp_avg_sq'],
8
    self._adamw_step_t, self._adamw_lr_t, self._adamw_beta1_t,
9
    self._adamw_beta2_t, self._adamw_eps_t, self._adamw_wd_t,
10
)

1
if not pinfo['is_small']:
2
        future = dist.all_gather_into_tensor(p, p_slice, async_op=True).get_future()