{"id":177748,"date":"2026-01-04T14:18:11","date_gmt":"2026-01-04T14:18:11","guid":{"rendered":"https:\/\/ktromedia.com\/?p=177748"},"modified":"2026-01-04T14:18:11","modified_gmt":"2026-01-04T14:18:11","slug":"train-your-large-model-on-multiple-gpus-with-pipeline-parallelism","status":"publish","type":"post","link":"http:\/\/ktromedia.com\/?p=177748","title":{"rendered":"Train Your Large Model on Multiple GPUs with Pipeline Parallelism"},"content":{"rendered":"
\n

import <\/span>dataclasses<\/span><\/p>\n

import <\/span>os<\/span><\/p>\n

\u00a0<\/p>\n

import <\/span>datasets<\/span><\/p>\n

import <\/span>tokenizers<\/span><\/p>\n

import <\/span>torch<\/span><\/p>\n

import <\/span>torch<\/span>.<\/span>distributed <\/span>as<\/span> <\/span>dist<\/span><\/p>\n

import <\/span>torch<\/span>.<\/span>nn <\/span>as<\/span> <\/span>nn<\/span><\/p>\n

import <\/span>torch<\/span>.<\/span>nn<\/span>.<\/span>functional <\/span>as<\/span> <\/span>F<\/span><\/p>\n

import <\/span>torch<\/span>.<\/span>optim<\/span>.<\/span>lr_scheduler <\/span>as<\/span> <\/span>lr_scheduler<\/span><\/p>\n

import <\/span>tqdm<\/span><\/p>\n

from <\/span>torch <\/span>import <\/span>Tensor<\/span><\/p>\n

from <\/span>torch<\/span>.<\/span>distributed<\/span>.<\/span>checkpoint <\/span>import <\/span>load<\/span>,<\/span> <\/span>save<\/span><\/p>\n

from <\/span>torch<\/span>.<\/span>distributed<\/span>.<\/span>checkpoint<\/span>.<\/span>state_dict <\/span>import <\/span>StateDictOptions<\/span>,<\/span> <\/span>get_state_dict<\/span>,<\/span> <\/span>set_state_dict<\/span><\/p>\n

from <\/span>torch<\/span>.<\/span>distributed<\/span>.<\/span>pipelining <\/span>import <\/span>PipelineStage<\/span>,<\/span> <\/span>ScheduleGPipe<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0<\/p>\n

# Build the model<\/span><\/p>\n

@<\/span>dataclasses<\/span>.<\/span>dataclass<\/span><\/p>\n

class<\/span> <\/span>LlamaConfig<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>“”<\/span>“Define Llama model hyperparameters.”<\/span>“”<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>vocab_size<\/span>:<\/span> <\/span>int<\/span> <\/span>=<\/span> <\/span>50000<\/span>\u00a0\u00a0<\/span># Size of the tokenizer vocabulary<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>max_position_embeddings<\/span>:<\/span> <\/span>int<\/span> <\/span>=<\/span> <\/span>2048<\/span>\u00a0\u00a0<\/span># Maximum sequence length<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>hidden_size<\/span>:<\/span> <\/span>int<\/span> <\/span>=<\/span> <\/span>768<\/span>\u00a0\u00a0<\/span># Dimension of hidden layers<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>intermediate_size<\/span>:<\/span> <\/span>int<\/span> <\/span>=<\/span> <\/span>4<\/span>*<\/span>768<\/span>\u00a0\u00a0<\/span># Dimension of MLP’s hidden layer<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>num_hidden_layers<\/span>:<\/span> <\/span>int<\/span> <\/span>=<\/span> <\/span>12<\/span>\u00a0\u00a0<\/span># Number of transformer layers<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>num_attention_heads<\/span>:<\/span> <\/span>int<\/span> <\/span>=<\/span> <\/span>12<\/span>\u00a0\u00a0<\/span># Number of attention heads<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>num_key_value_heads<\/span>:<\/span> <\/span>int<\/span> <\/span>=<\/span> <\/span>3<\/span>\u00a0\u00a0<\/span># Number of key-value heads for GQA<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0<\/p>\n

class<\/span> <\/span>RotaryPositionEncoding<\/span>(<\/span>nn<\/span>.<\/span>Module<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>“”<\/span>“Rotary position encoding.”<\/span>“”<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>def <\/span>__init__<\/span>(<\/span>self<\/span>,<\/span> <\/span>dim<\/span>:<\/span> <\/span>int<\/span>,<\/span> <\/span>max_position_embeddings<\/span>:<\/span> <\/span>int<\/span>)<\/span> <\/span>-><\/span> <\/span>None<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>“”<\/span>“Initialize the RotaryPositionEncoding module.<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0Args:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0dim: The hidden dimension of the input tensor to which RoPE is applied<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0max_position_embeddings: The maximum sequence length of the input tensor<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0“<\/span>“”<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>super<\/span>(<\/span>)<\/span>.<\/span>__init__<\/span>(<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>dim<\/span> <\/span>=<\/span> <\/span>dim<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>max_position_embeddings<\/span> <\/span>=<\/span> <\/span>max_position<\/span>_<\/span>embeddings<\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># compute a matrix of n\\theta_i<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>N<\/span> <\/span>=<\/span> <\/span>10_000.0<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>inv_freq<\/span> <\/span>=<\/span> <\/span>1.0<\/span> <\/span>\/<\/span> <\/span>(<\/span>N *<\/span>*<\/span> <\/span>(<\/span>torch<\/span>.<\/span>arange<\/span>(<\/span>0<\/span>,<\/span> <\/span>dim<\/span>,<\/span> <\/span>2<\/span>)<\/span> <\/span>\/<\/span> <\/span>dim<\/span>)<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>inv_freq<\/span> <\/span>=<\/span> <\/span>torch<\/span>.<\/span>cat<\/span>(<\/span>(<\/span>inv_freq<\/span>,<\/span> <\/span>inv_freq<\/span>)<\/span>,<\/span> <\/span>dim<\/span>=<\/span>–<\/span>1<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>position<\/span> <\/span>=<\/span> <\/span>torch<\/span>.<\/span>arange<\/span>(<\/span>max_position_embeddings<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>sinusoid_inp<\/span> <\/span>=<\/span> <\/span>torch<\/span>.<\/span>outer<\/span>(<\/span>position<\/span>,<\/span> <\/span>inv_freq<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># save cosine and sine matrices as buffers, not parameters<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>register_buffer<\/span>(<\/span>“cos”<\/span>,<\/span> <\/span>sinusoid_inp<\/span>.<\/span>cos<\/span>(<\/span>)<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>register_buffer<\/span>(<\/span>“sin”<\/span>,<\/span> <\/span>sinusoid_inp<\/span>.<\/span>sin<\/span>(<\/span>)<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>def <\/span>forward<\/span>(<\/span>self<\/span>,<\/span> <\/span>x<\/span>:<\/span> <\/span>Tensor<\/span>)<\/span> <\/span>-><\/span> <\/span>Tensor<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>“”<\/span>“Apply RoPE to tensor x.<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0Args:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0x: Input tensor of shape (batch_size, seq_length, num_heads, head_dim)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0Returns:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0Output tensor of shape (batch_size, seq_length, num_heads, head_dim)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0“<\/span>“”<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>batch_size<\/span>,<\/span> <\/span>seq_len<\/span>,<\/span> <\/span>num_heads<\/span>,<\/span> <\/span>head_dim<\/span> <\/span>=<\/span> <\/span>x<\/span>.<\/span>shape<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>dtype<\/span> <\/span>=<\/span> <\/span>x<\/span>.<\/span>dtype<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># transform the cosine and sine matrices to 4D tensor and the same dtype as x<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>cos<\/span> <\/span>=<\/span> <\/span>self<\/span>.<\/span>cos<\/span>.<\/span>to<\/span>(<\/span>dtype<\/span>)<\/span>[<\/span>:<\/span>seq_len<\/span>]<\/span>.<\/span>view<\/span>(<\/span>1<\/span>,<\/span> <\/span>seq_len<\/span>,<\/span> <\/span>1<\/span>,<\/span> <\/span>–<\/span>1<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>sin<\/span> <\/span>=<\/span> <\/span>self<\/span>.<\/span>sin<\/span>.<\/span>to<\/span>(<\/span>dtype<\/span>)<\/span>[<\/span>:<\/span>seq_len<\/span>]<\/span>.<\/span>view<\/span>(<\/span>1<\/span>,<\/span> <\/span>seq_len<\/span>,<\/span> <\/span>1<\/span>,<\/span> <\/span>–<\/span>1<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># apply RoPE to x<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>x1<\/span>,<\/span> <\/span>x2<\/span> <\/span>=<\/span> <\/span>x<\/span>.<\/span>chunk<\/span>(<\/span>2<\/span>,<\/span> <\/span>dim<\/span>=<\/span>–<\/span>1<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>rotated<\/span> <\/span>=<\/span> <\/span>torch<\/span>.<\/span>cat<\/span>(<\/span>(<\/span>–<\/span>x2<\/span>,<\/span> <\/span>x1<\/span>)<\/span>,<\/span> <\/span>dim<\/span>=<\/span>–<\/span>1<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>output<\/span> <\/span>=<\/span> <\/span>(<\/span>x *<\/span> <\/span>cos<\/span>)<\/span> <\/span>+<\/span> <\/span>(<\/span>rotated *<\/span> <\/span>sin<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>return<\/span> <\/span>output<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0<\/p>\n

class<\/span> <\/span>LlamaAttention<\/span>(<\/span>nn<\/span>.<\/span>Module<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>“”<\/span>“Grouped-query attention with rotary embeddings.”<\/span>“”<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>def <\/span>__init__<\/span>(<\/span>self<\/span>,<\/span> <\/span>config<\/span>:<\/span> <\/span>LlamaConfig<\/span>)<\/span> <\/span>-><\/span> <\/span>None<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>super<\/span>(<\/span>)<\/span>.<\/span>__init__<\/span>(<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>hidden_size<\/span> <\/span>=<\/span> <\/span>config<\/span>.<\/span>hidden_size<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>num_heads<\/span> <\/span>=<\/span> <\/span>config<\/span>.<\/span>num_attention_heads<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>head_dim<\/span> <\/span>=<\/span> <\/span>self<\/span>.<\/span>hidden_size<\/span> <\/span>\/\/ self.num_heads<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>num_kv_heads<\/span> <\/span>=<\/span> <\/span>config<\/span>.<\/span>num_key_value<\/span>_<\/span>heads\u00a0\u00a0<\/span># GQA: H_kv <\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># hidden_size must be divisible by num_heads<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>assert<\/span> <\/span>(<\/span>self<\/span>.<\/span>head_dim *<\/span> <\/span>self<\/span>.<\/span>num_heads<\/span>)<\/span> <\/span>==<\/span> <\/span>self<\/span>.<\/span>hidden<\/span>_<\/span>size<\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># Linear layers for Q, K, V projections<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>q_proj<\/span> <\/span>=<\/span> <\/span>nn<\/span>.<\/span>Linear<\/span>(<\/span>self<\/span>.<\/span>hidden_size<\/span>,<\/span> <\/span>self<\/span>.<\/span>num_heads *<\/span> <\/span>self<\/span>.<\/span>head_dim<\/span>,<\/span> <\/span>bias<\/span>=<\/span>False<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>k_proj<\/span> <\/span>=<\/span> <\/span>nn<\/span>.<\/span>Linear<\/span>(<\/span>self<\/span>.<\/span>hidden_size<\/span>,<\/span> <\/span>self<\/span>.<\/span>num_kv_heads *<\/span> <\/span>self<\/span>.<\/span>head_dim<\/span>,<\/span> <\/span>bias<\/span>=<\/span>False<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>v_proj<\/span> <\/span>=<\/span> <\/span>nn<\/span>.<\/span>Linear<\/span>(<\/span>self<\/span>.<\/span>hidden_size<\/span>,<\/span> <\/span>self<\/span>.<\/span>num_kv_heads *<\/span> <\/span>self<\/span>.<\/span>head_dim<\/span>,<\/span> <\/span>bias<\/span>=<\/span>False<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>o_proj<\/span> <\/span>=<\/span> <\/span>nn<\/span>.<\/span>Linear<\/span>(<\/span>self<\/span>.<\/span>num_heads *<\/span> <\/span>self<\/span>.<\/span>head_dim<\/span>,<\/span> <\/span>self<\/span>.<\/span>hidden_size<\/span>,<\/span> <\/span>bias<\/span>=<\/span>False<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>def <\/span>forward<\/span>(<\/span>self<\/span>,<\/span> <\/span>hidden_states<\/span>:<\/span> <\/span>Tensor<\/span>,<\/span> <\/span>rope<\/span>:<\/span> <\/span>RotaryPositionEncoding<\/span>)<\/span> <\/span>-><\/span> <\/span>Tensor<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>bs<\/span>,<\/span> <\/span>seq_len<\/span>,<\/span> <\/span>dim<\/span> <\/span>=<\/span> <\/span>hidden_states<\/span>.<\/span>size<\/span>(<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># Project inputs to Q, K, V<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>query_states<\/span> <\/span>=<\/span> <\/span>self<\/span>.<\/span>q_proj<\/span>(<\/span>hidden_states<\/span>)<\/span>.<\/span>view<\/span>(<\/span>bs<\/span>,<\/span> <\/span>seq_len<\/span>,<\/span> <\/span>self<\/span>.<\/span>num_heads<\/span>,<\/span> <\/span>self<\/span>.<\/span>head_dim<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>key_states<\/span> <\/span>=<\/span> <\/span>self<\/span>.<\/span>k_proj<\/span>(<\/span>hidden_states<\/span>)<\/span>.<\/span>view<\/span>(<\/span>bs<\/span>,<\/span> <\/span>seq_len<\/span>,<\/span> <\/span>self<\/span>.<\/span>num_kv_heads<\/span>,<\/span> <\/span>self<\/span>.<\/span>head_dim<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>value_states<\/span> <\/span>=<\/span> <\/span>self<\/span>.<\/span>v_proj<\/span>(<\/span>hidden_states<\/span>)<\/span>.<\/span>view<\/span>(<\/span>bs<\/span>,<\/span> <\/span>seq_len<\/span>,<\/span> <\/span>self<\/span>.<\/span>num_kv_heads<\/span>,<\/span> <\/span>self<\/span>.<\/span>head_dim<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># Apply rotary position embeddings<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>query_states<\/span> <\/span>=<\/span> <\/span>rope<\/span>(<\/span>query_states<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>key_states<\/span> <\/span>=<\/span> <\/span>rope<\/span>(<\/span>key_states<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># Transpose tensors from BSHD to BHSD dimension for scaled_dot_product_attention<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>query_states<\/span> <\/span>=<\/span> <\/span>query_states<\/span>.<\/span>transpose<\/span>(<\/span>1<\/span>,<\/span> <\/span>2<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>key_states<\/span> <\/span>=<\/span> <\/span>key_states<\/span>.<\/span>transpose<\/span>(<\/span>1<\/span>,<\/span> <\/span>2<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>value_states<\/span> <\/span>=<\/span> <\/span>value_states<\/span>.<\/span>transpose<\/span>(<\/span>1<\/span>,<\/span> <\/span>2<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># Use PyTorch’s optimized attention implementation<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># setting is_causal=True is incompatible with setting explicit attention mask<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>attn_output<\/span> <\/span>=<\/span> <\/span>F<\/span>.<\/span>scaled_dot_product_attention<\/span>(<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>query_states<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>key_states<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>value_states<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>is_causal<\/span>=<\/span>True<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>dropout_p<\/span>=<\/span>0.0<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>enable_gqa<\/span>=<\/span>True<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># Transpose output tensor from BHSD to BSHD dimension, reshape to 3D, and then project output<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>attn_output<\/span> <\/span>=<\/span> <\/span>attn_output<\/span>.<\/span>transpose<\/span>(<\/span>1<\/span>,<\/span> <\/span>2<\/span>)<\/span>.<\/span>reshape<\/span>(<\/span>bs<\/span>,<\/span> <\/span>seq_len<\/span>,<\/span> <\/span>self<\/span>.<\/span>hidden_size<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>attn_output<\/span> <\/span>=<\/span> <\/span>self<\/span>.<\/span>o_proj<\/span>(<\/span>attn_output<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>return<\/span> <\/span>attn_output<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0<\/p>\n

class<\/span> <\/span>LlamaMLP<\/span>(<\/span>nn<\/span>.<\/span>Module<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>“”<\/span>“Feed-forward network with SwiGLU activation.”<\/span>“”<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>def <\/span>__init__<\/span>(<\/span>self<\/span>,<\/span> <\/span>config<\/span>:<\/span> <\/span>LlamaConfig<\/span>)<\/span> <\/span>-><\/span> <\/span>None<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>super<\/span>(<\/span>)<\/span>.<\/span>__init__<\/span>(<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># Two parallel projections for SwiGLU<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>gate_proj<\/span> <\/span>=<\/span> <\/span>nn<\/span>.<\/span>Linear<\/span>(<\/span>config<\/span>.<\/span>hidden_size<\/span>,<\/span> <\/span>config<\/span>.<\/span>intermediate_size<\/span>,<\/span> <\/span>bias<\/span>=<\/span>False<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>up_proj<\/span> <\/span>=<\/span> <\/span>nn<\/span>.<\/span>Linear<\/span>(<\/span>config<\/span>.<\/span>hidden_size<\/span>,<\/span> <\/span>config<\/span>.<\/span>intermediate_size<\/span>,<\/span> <\/span>bias<\/span>=<\/span>False<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>act_fn<\/span> <\/span>=<\/span> <\/span>F<\/span>.<\/span>silu<\/span>\u00a0\u00a0<\/span># SwiGLU activation function<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># Project back to hidden size<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>down_proj<\/span> <\/span>=<\/span> <\/span>nn<\/span>.<\/span>Linear<\/span>(<\/span>config<\/span>.<\/span>intermediate_size<\/span>,<\/span> <\/span>config<\/span>.<\/span>hidden_size<\/span>,<\/span> <\/span>bias<\/span>=<\/span>False<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>def <\/span>forward<\/span>(<\/span>self<\/span>,<\/span> <\/span>x<\/span>:<\/span> <\/span>Tensor<\/span>)<\/span> <\/span>-><\/span> <\/span>Tensor<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># SwiGLU activation: multiply gate and up-projected inputs<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>gate<\/span> <\/span>=<\/span> <\/span>self<\/span>.<\/span>act_fn<\/span>(<\/span>self<\/span>.<\/span>gate_proj<\/span>(<\/span>x<\/span>)<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>up<\/span> <\/span>=<\/span> <\/span>self<\/span>.<\/span>up_proj<\/span>(<\/span>x<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>return<\/span> <\/span>self<\/span>.<\/span>down_proj<\/span>(<\/span>gate *<\/span> <\/span>up<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0<\/p>\n

class<\/span> <\/span>LlamaDecoderLayer<\/span>(<\/span>nn<\/span>.<\/span>Module<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>“”<\/span>“Single transformer layer for a Llama model.”<\/span>“”<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>def <\/span>__init__<\/span>(<\/span>self<\/span>,<\/span> <\/span>config<\/span>:<\/span> <\/span>LlamaConfig<\/span>)<\/span> <\/span>-><\/span> <\/span>None<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>super<\/span>(<\/span>)<\/span>.<\/span>__init__<\/span>(<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>input_layernorm<\/span> <\/span>=<\/span> <\/span>nn<\/span>.<\/span>RMSNorm<\/span>(<\/span>config<\/span>.<\/span>hidden_size<\/span>,<\/span> <\/span>eps<\/span>=<\/span>1e<\/span>–<\/span>5<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>self_attn<\/span> <\/span>=<\/span> <\/span>LlamaAttention<\/span>(<\/span>config<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>post_attention_layernorm<\/span> <\/span>=<\/span> <\/span>nn<\/span>.<\/span>RMSNorm<\/span>(<\/span>config<\/span>.<\/span>hidden_size<\/span>,<\/span> <\/span>eps<\/span>=<\/span>1e<\/span>–<\/span>5<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>mlp<\/span> <\/span>=<\/span> <\/span>LlamaMLP<\/span>(<\/span>config<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>def <\/span>forward<\/span>(<\/span>self<\/span>,<\/span> <\/span>hidden_states<\/span>:<\/span> <\/span>Tensor<\/span>,<\/span> <\/span>rope<\/span>:<\/span> <\/span>RotaryPositionEncoding<\/span>)<\/span> <\/span>-><\/span> <\/span>Tensor<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># First residual block: Self-attention<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>residual<\/span> <\/span>=<\/span> <\/span>hidden_states<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>hidden_states<\/span> <\/span>=<\/span> <\/span>self<\/span>.<\/span>input_layernorm<\/span>(<\/span>hidden_states<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>attn_outputs<\/span> <\/span>=<\/span> <\/span>self<\/span>.<\/span>self_attn<\/span>(<\/span>hidden_states<\/span>,<\/span> <\/span>rope<\/span>=<\/span>rope<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>hidden_states<\/span> <\/span>=<\/span> <\/span>attn_outputs<\/span> <\/span>+<\/span> <\/span>residual<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># Second residual block: MLP<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>residual<\/span> <\/span>=<\/span> <\/span>hidden_states<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>hidden_states<\/span> <\/span>=<\/span> <\/span>self<\/span>.<\/span>post_attention_layernorm<\/span>(<\/span>hidden_states<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>hidden_states<\/span> <\/span>=<\/span> <\/span>self<\/span>.<\/span>mlp<\/span>(<\/span>hidden_states<\/span>)<\/span> <\/span>+<\/span> <\/span>residual<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>return<\/span> <\/span>hidden_states<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0<\/p>\n

class<\/span> <\/span>LlamaModel<\/span>(<\/span>nn<\/span>.<\/span>Module<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>“”<\/span>“The full Llama model without any pretraining heads.”<\/span>“”<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>def <\/span>__init__<\/span>(<\/span>self<\/span>,<\/span> <\/span>config<\/span>:<\/span> <\/span>LlamaConfig<\/span>)<\/span> <\/span>-><\/span> <\/span>None<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>super<\/span>(<\/span>)<\/span>.<\/span>__init__<\/span>(<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>rope<\/span> <\/span>=<\/span> <\/span>RotaryPositionEncoding<\/span>(<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>config<\/span>.<\/span>hidden_size<\/span> <\/span>\/\/ config.num_attention_heads,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>config<\/span>.<\/span>max_position_embeddings<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>embed_tokens<\/span> <\/span>=<\/span> <\/span>nn<\/span>.<\/span>Embedding<\/span>(<\/span>config<\/span>.<\/span>vocab_size<\/span>,<\/span> <\/span>config<\/span>.<\/span>hidden_size<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>layers<\/span> <\/span>=<\/span> <\/span>nn<\/span>.<\/span>ModuleDict<\/span>(<\/span>{<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>str<\/span>(<\/span>i<\/span>)<\/span>:<\/span> <\/span>LlamaDecoderLayer<\/span>(<\/span>config<\/span>)<\/span> <\/span>for<\/span> <\/span>i<\/span> <\/span>in<\/span> <\/span>range<\/span>(<\/span>config<\/span>.<\/span>num_hidden_layers<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>}<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>norm<\/span> <\/span>=<\/span> <\/span>nn<\/span>.<\/span>RMSNorm<\/span>(<\/span>config<\/span>.<\/span>hidden_size<\/span>,<\/span> <\/span>eps<\/span>=<\/span>1e<\/span>–<\/span>5<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>def <\/span>forward<\/span>(<\/span>self<\/span>,<\/span> <\/span>input_ids<\/span>:<\/span> <\/span>Tensor<\/span>)<\/span> <\/span>-><\/span> <\/span>Tensor<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># Convert input token IDs to embeddings<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>if<\/span> <\/span>self<\/span>.<\/span>embed_tokens <\/span>is<\/span> <\/span>not<\/span> <\/span>None<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>hidden_states<\/span> <\/span>=<\/span> <\/span>self<\/span>.<\/span>embed_tokens<\/span>(<\/span>input_ids<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>else<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>hidden_states<\/span> <\/span>=<\/span> <\/span>input<\/span>_<\/span>ids<\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># Process through all transformer layers, then the final norm layer<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>for<\/span> <\/span>n<\/span> <\/span>in<\/span> <\/span>range<\/span>(<\/span>len<\/span>(<\/span>self<\/span>.<\/span>layers<\/span>)<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>if<\/span> <\/span>self<\/span>.<\/span>layers<\/span>[<\/span>str<\/span>(<\/span>n<\/span>)<\/span>]<\/span> <\/span>is<\/span> <\/span>not<\/span> <\/span>None<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>hidden_states<\/span> <\/span>=<\/span> <\/span>self<\/span>.<\/span>layers<\/span>[<\/span>str<\/span>(<\/span>n<\/span>)<\/span>]<\/span>(<\/span>hidden_states<\/span>,<\/span> <\/span>self<\/span>.<\/span>rope<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>if<\/span> <\/span>self<\/span>.<\/span>norm <\/span>is<\/span> <\/span>not<\/span> <\/span>None<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>hidden_states<\/span> <\/span>=<\/span> <\/span>self<\/span>.<\/span>norm<\/span>(<\/span>hidden_states<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># Return the final hidden states, and copy over the attention mask<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>return<\/span> <\/span>hidden_states<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0<\/p>\n

class<\/span> <\/span>LlamaForPretraining<\/span>(<\/span>nn<\/span>.<\/span>Module<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>def <\/span>__init__<\/span>(<\/span>self<\/span>,<\/span> <\/span>config<\/span>:<\/span> <\/span>LlamaConfig<\/span>)<\/span> <\/span>-><\/span> <\/span>None<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>super<\/span>(<\/span>)<\/span>.<\/span>__init__<\/span>(<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>base_model<\/span> <\/span>=<\/span> <\/span>LlamaModel<\/span>(<\/span>config<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>lm_head<\/span> <\/span>=<\/span> <\/span>nn<\/span>.<\/span>Linear<\/span>(<\/span>config<\/span>.<\/span>hidden_size<\/span>,<\/span> <\/span>config<\/span>.<\/span>vocab_size<\/span>,<\/span> <\/span>bias<\/span>=<\/span>False<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>def <\/span>forward<\/span>(<\/span>self<\/span>,<\/span> <\/span>input_ids<\/span>:<\/span> <\/span>Tensor<\/span>)<\/span> <\/span>-><\/span> <\/span>Tensor<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>hidden_states<\/span> <\/span>=<\/span> <\/span>self<\/span>.<\/span>base_model<\/span>(<\/span>input_ids<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>if<\/span> <\/span>self<\/span>.<\/span>lm_head <\/span>is<\/span> <\/span>not<\/span> <\/span>None<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>hidden_states<\/span> <\/span>=<\/span> <\/span>self<\/span>.<\/span>lm_head<\/span>(<\/span>hidden_states<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>return<\/span> <\/span>hidden<\/span>_<\/span>states<\/p>\n

\u00a0<\/p>\n

\u00a0<\/p>\n

# Generator function to create padded sequences of fixed length<\/span><\/p>\n

class<\/span> <\/span>PretrainingDataset<\/span>(<\/span>torch<\/span>.<\/span>utils<\/span>.<\/span>data<\/span>.<\/span>Dataset<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>def <\/span>__init__<\/span>(<\/span>self<\/span>,<\/span> <\/span>dataset<\/span>:<\/span> <\/span>datasets<\/span>.<\/span>Dataset<\/span>,<\/span> <\/span>tokenizer<\/span>:<\/span> <\/span>tokenizers<\/span>.<\/span>Tokenizer<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0 <\/span>seq_length<\/span>:<\/span> <\/span>int<\/span>,<\/span> <\/span>device<\/span>:<\/span> <\/span>torch<\/span>.<\/span>device<\/span> <\/span>=<\/span> <\/span>None<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>dataset<\/span> <\/span>=<\/span> <\/span>dataset<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>tokenizer<\/span> <\/span>=<\/span> <\/span>tokenizer<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>device<\/span> <\/span>=<\/span> <\/span>device<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>seq_length<\/span> <\/span>=<\/span> <\/span>seq_length<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>bot<\/span> <\/span>=<\/span> <\/span>tokenizer<\/span>.<\/span>token_to_id<\/span>(<\/span>“[BOT]”<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>eot<\/span> <\/span>=<\/span> <\/span>tokenizer<\/span>.<\/span>token_to_id<\/span>(<\/span>“[EOT]”<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>self<\/span>.<\/span>pad<\/span> <\/span>=<\/span> <\/span>tokenizer<\/span>.<\/span>token_to_id<\/span>(<\/span>“[PAD]”<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>def <\/span>__len__<\/span>(<\/span>self<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>return<\/span> <\/span>len<\/span>(<\/span>self<\/span>.<\/span>dataset<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>def <\/span>__getitem__<\/span>(<\/span>self<\/span>,<\/span> <\/span>index<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>“”<\/span>“Get a sequence of token ids from the dataset. [BOT] and [EOT] tokens<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0are added. Clipped and padded to the sequence length.<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0“<\/span>“”<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>seq<\/span> <\/span>=<\/span> <\/span>self<\/span>.<\/span>dataset<\/span>[<\/span>index<\/span>]<\/span>[<\/span>“text”<\/span>]<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>tokens<\/span>:<\/span> <\/span>list<\/span>[<\/span>int<\/span>]<\/span> <\/span>=<\/span> <\/span>[<\/span>self<\/span>.<\/span>bot<\/span>]<\/span> <\/span>+<\/span> <\/span>self<\/span>.<\/span>tokenizer<\/span>.<\/span>encode<\/span>(<\/span>seq<\/span>)<\/span>.<\/span>ids<\/span> <\/span>+<\/span> <\/span>[<\/span>self<\/span>.<\/span>eot<\/span>]<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># pad to target sequence length<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>toklen<\/span> <\/span>=<\/span> <\/span>len<\/span>(<\/span>tokens<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>if<\/span> <\/span>toklen<\/span> <\/span> <\/span>self<\/span>.<\/span>seq_length<\/span>+<\/span>1<\/span>:<\/span><\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>pad_length<\/span> <\/span>=<\/span> <\/span>self<\/span>.<\/span>seq_length<\/span>+<\/span>1<\/span> <\/span>–<\/span> <\/span>toklen<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>tokens<\/span> <\/span>+=<\/span> <\/span>[<\/span>self<\/span>.<\/span>pad<\/span>]<\/span> <\/span>*<\/span> <\/span>pad<\/span>_<\/span>length<\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># return the sequence<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>x<\/span> <\/span>=<\/span> <\/span>torch<\/span>.<\/span>tensor<\/span>(<\/span>tokens<\/span>[<\/span>:<\/span>self<\/span>.<\/span>seq_length<\/span>]<\/span>,<\/span> <\/span>dtype<\/span>=<\/span>torch<\/span>.<\/span>int64<\/span>,<\/span> <\/span>device<\/span>=<\/span>self<\/span>.<\/span>device<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>y<\/span> <\/span>=<\/span> <\/span>torch<\/span>.<\/span>tensor<\/span>(<\/span>tokens<\/span>[<\/span>1<\/span>:<\/span>self<\/span>.<\/span>seq_length<\/span>+<\/span>1<\/span>]<\/span>,<\/span> <\/span>dtype<\/span>=<\/span>torch<\/span>.<\/span>int64<\/span>,<\/span> <\/span>device<\/span>=<\/span>self<\/span>.<\/span>device<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>return<\/span> <\/span>x<\/span>,<\/span> <\/span>y<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0<\/p>\n

def <\/span>load_checkpoint<\/span>(<\/span>model<\/span>:<\/span> <\/span>nn<\/span>.<\/span>Module<\/span>,<\/span> <\/span>optimizer<\/span>:<\/span> <\/span>torch<\/span>.<\/span>optim<\/span>.<\/span>Optimizer<\/span>)<\/span> <\/span>-><\/span> <\/span>None<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>dist<\/span>.<\/span>barrier<\/span>(<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>model_state<\/span>,<\/span> <\/span>optimizer_state<\/span> <\/span>=<\/span> <\/span>get_state_dict<\/span>(<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>model<\/span>,<\/span> <\/span>optimizer<\/span>,<\/span> <\/span>options<\/span>=<\/span>StateDictOptions<\/span>(<\/span>full_state_dict<\/span>=<\/span>True<\/span>)<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>load<\/span>(<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>{<\/span>“model”<\/span>:<\/span> <\/span>model_state<\/span>,<\/span> <\/span>“optimizer”<\/span>:<\/span> <\/span>optimizer_state<\/span>}<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>checkpoint_id<\/span>=<\/span>“checkpoint-dist”<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>set_state_dict<\/span>(<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>model<\/span>,<\/span> <\/span>optimizer<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>model_state_dict<\/span>=<\/span>model_state<\/span>,<\/span> <\/span>optim_state_dict<\/span>=<\/span>optimizer_state<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>options<\/span>=<\/span>StateDictOptions<\/span>(<\/span>broadcast_from_rank0<\/span>=<\/span>True<\/span>,<\/span> <\/span>full_state_dict<\/span>=<\/span>True<\/span>)<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>dist<\/span>.<\/span>barrier<\/span>(<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0<\/p>\n

def <\/span>save_checkpoint<\/span>(<\/span>model<\/span>:<\/span> <\/span>nn<\/span>.<\/span>Module<\/span>,<\/span> <\/span>optimizer<\/span>:<\/span> <\/span>torch<\/span>.<\/span>optim<\/span>.<\/span>Optimizer<\/span>)<\/span> <\/span>-><\/span> <\/span>None<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>dist<\/span>.<\/span>barrier<\/span>(<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>model_state<\/span>,<\/span> <\/span>optimizer_state<\/span> <\/span>=<\/span> <\/span>get_state_dict<\/span>(<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>model<\/span>,<\/span> <\/span>optimizer<\/span>,<\/span> <\/span>options<\/span>=<\/span>StateDictOptions<\/span>(<\/span>full_state_dict<\/span>=<\/span>True<\/span>)<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>save<\/span>(<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>{<\/span>“model”<\/span>:<\/span> <\/span>model_state<\/span>,<\/span> <\/span>“optimizer”<\/span>:<\/span> <\/span>optimizer_state<\/span>}<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>checkpoint_id<\/span>=<\/span>“checkpoint-dist”<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>dist<\/span>.<\/span>barrier<\/span>(<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0<\/p>\n

# Load the tokenizer and dataset<\/span><\/p>\n

tokenizer<\/span> <\/span>=<\/span> <\/span>tokenizers<\/span>.<\/span>Tokenizer<\/span>.<\/span>from_file<\/span>(<\/span>“bpe_50K.json”<\/span>)<\/span><\/p>\n

dataset<\/span> <\/span>=<\/span> <\/span>datasets<\/span>.<\/span>load_dataset<\/span>(<\/span>“HuggingFaceFW\/fineweb”<\/span>,<\/span> <\/span>“sample-10BT”<\/span>,<\/span> <\/span>split<\/span>=<\/span>“train”<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

# Initialize the distributed environment<\/span><\/p>\n

dist<\/span>.<\/span>init_process_group<\/span>(<\/span>backend<\/span>=<\/span>“nccl”<\/span>)<\/span><\/p>\n

rank<\/span> <\/span>=<\/span> <\/span>dist<\/span>.<\/span>get_rank<\/span>(<\/span>)<\/span><\/p>\n

local_rank<\/span> <\/span>=<\/span> <\/span>int<\/span>(<\/span>os<\/span>.<\/span>environ<\/span>[<\/span>“LOCAL_RANK”<\/span>]<\/span>)<\/span><\/p>\n

world_size<\/span> <\/span>=<\/span> <\/span>dist<\/span>.<\/span>get_world_size<\/span>(<\/span>)<\/span><\/p>\n

device<\/span> <\/span>=<\/span> <\/span>torch<\/span>.<\/span>device<\/span>(<\/span>f<\/span>“cuda:{local_rank}”<\/span>)<\/span><\/p>\n

print<\/span>(<\/span>f<\/span>“World size {world_size}, rank {rank}, local rank {local_rank}. Using {device}”<\/span>)<\/span><\/p>\n

assert<\/span> <\/span>world_size<\/span> <\/span>==<\/span> <\/span>3<\/span>,<\/span> <\/span>f<\/span>“This script is designed for 3 GPUs, got {world_size}”<\/span><\/p>\n

\u00a0<\/p>\n

# Create pretraining model with default config on meta device to prevent OOM<\/span><\/p>\n

with <\/span>torch<\/span>.<\/span>device<\/span>(<\/span>“meta”<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>model_config<\/span> <\/span>=<\/span> <\/span>LlamaConfig<\/span>(<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>model<\/span> <\/span>=<\/span> <\/span>LlamaForPretraining<\/span>(<\/span>model_config<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span># Partition the model by removing some layers<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>num_layers<\/span> <\/span>=<\/span> <\/span>model_config<\/span>.<\/span>num_hidden_layers<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>partition<\/span> <\/span>=<\/span> <\/span>[<\/span>num_layers<\/span> <\/span>\/\/ 3, 2 * num_layers \/\/ 3, num_layers]<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>if<\/span> <\/span>rank<\/span> <\/span>==<\/span> <\/span>0<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># from embedding to 1\/3 of the decoder layers<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>for<\/span> <\/span>n<\/span> <\/span>in<\/span> <\/span>range<\/span>(<\/span>partition<\/span>[<\/span>0<\/span>]<\/span>,<\/span> <\/span>partition<\/span>[<\/span>2<\/span>]<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>model<\/span>.<\/span>base_model<\/span>.<\/span>layers<\/span>[<\/span>str<\/span>(<\/span>n<\/span>)<\/span>]<\/span> <\/span>=<\/span> <\/span>None<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>model<\/span>.<\/span>base_model<\/span>.<\/span>norm<\/span> <\/span>=<\/span> <\/span>None<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>model<\/span>.<\/span>lm_head<\/span> <\/span>=<\/span> <\/span>None<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>elif <\/span>rank<\/span> <\/span>==<\/span> <\/span>1<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># from 1\/3 to 2\/3 of the decoder layers<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>model<\/span>.<\/span>base_model<\/span>.<\/span>embed_tokens<\/span> <\/span>=<\/span> <\/span>None<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>for<\/span> <\/span>n<\/span> <\/span>in<\/span> <\/span>range<\/span>(<\/span>0<\/span>,<\/span> <\/span>partition<\/span>[<\/span>0<\/span>]<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>model<\/span>.<\/span>base_model<\/span>.<\/span>layers<\/span>[<\/span>str<\/span>(<\/span>n<\/span>)<\/span>]<\/span> <\/span>=<\/span> <\/span>None<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>for<\/span> <\/span>n<\/span> <\/span>in<\/span> <\/span>range<\/span>(<\/span>partition<\/span>[<\/span>1<\/span>]<\/span>,<\/span> <\/span>partition<\/span>[<\/span>2<\/span>]<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>model<\/span>.<\/span>base_model<\/span>.<\/span>layers<\/span>[<\/span>str<\/span>(<\/span>n<\/span>)<\/span>]<\/span> <\/span>=<\/span> <\/span>None<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>model<\/span>.<\/span>base_model<\/span>.<\/span>norm<\/span> <\/span>=<\/span> <\/span>None<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>model<\/span>.<\/span>lm_head<\/span> <\/span>=<\/span> <\/span>None<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>elif <\/span>rank<\/span> <\/span>==<\/span> <\/span>2<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># from 2\/3 to the end of the decoder layers and the final norm layer, LM head<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>model<\/span>.<\/span>base_model<\/span>.<\/span>embed_tokens<\/span> <\/span>=<\/span> <\/span>None<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>for<\/span> <\/span>n<\/span> <\/span>in<\/span> <\/span>range<\/span>(<\/span>partition<\/span>[<\/span>1<\/span>]<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>model<\/span>.<\/span>base_model<\/span>.<\/span>layers<\/span>[<\/span>str<\/span>(<\/span>n<\/span>)<\/span>]<\/span> <\/span>=<\/span> <\/span>None<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>else<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>raise <\/span>ValueError<\/span>(<\/span>f<\/span>“Invalid rank: {rank}”<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0<\/p>\n

# Move model from meta device to CUDA device, then initialize the weights<\/span><\/p>\n

def <\/span>reset_all_weights<\/span>(<\/span>model<\/span>:<\/span> <\/span>nn<\/span>.<\/span>Module<\/span>)<\/span> <\/span>-><\/span> <\/span>None<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>@<\/span>torch<\/span>.<\/span>no_grad<\/span>(<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>def <\/span>weight_reset<\/span>(<\/span>m<\/span>:<\/span> <\/span>nn<\/span>.<\/span>Module<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>reset_parameters<\/span> <\/span>=<\/span> <\/span>getattr<\/span>(<\/span>m<\/span>,<\/span> <\/span>“reset_parameters”<\/span>,<\/span> <\/span>None<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>if<\/span> <\/span>callable<\/span>(<\/span>reset_parameters<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>m<\/span>.<\/span>reset_parameters<\/span>(<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span># Applies fn recursively to model itself and all of model.children()<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>model<\/span>.<\/span>apply<\/span>(<\/span>fn<\/span>=<\/span>weight_reset<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0<\/p>\n

model<\/span>.<\/span>to_empty<\/span>(<\/span>device<\/span>=<\/span>device<\/span>)<\/span><\/p>\n

reset_all_weights<\/span>(<\/span>model<\/span>)<\/span><\/p>\n

model<\/span>.<\/span>train<\/span>(<\/span>)<\/span><\/p>\n

stage<\/span> <\/span>=<\/span> <\/span>PipelineStage<\/span>(<\/span>model<\/span>,<\/span> <\/span>stage_index<\/span>=<\/span>rank<\/span>,<\/span> <\/span>num_stages<\/span>=<\/span>world_size<\/span>,<\/span> <\/span>device<\/span>=<\/span>device<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

# Training parameters<\/span><\/p>\n

epochs<\/span> <\/span>=<\/span> <\/span>3<\/span><\/p>\n

learning_rate<\/span> <\/span>=<\/span> <\/span>1e<\/span>–<\/span>3<\/span><\/p>\n

batch_size<\/span> <\/span>=<\/span> <\/span>64<\/span><\/p>\n

seq_length<\/span> <\/span>=<\/span> <\/span>512<\/span><\/p>\n

num_warmup_steps<\/span> <\/span>=<\/span> <\/span>1000<\/span><\/p>\n

PAD_TOKEN_ID<\/span> <\/span>=<\/span> <\/span>tokenizer<\/span>.<\/span>token_to_id<\/span>(<\/span>“[PAD]”<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

# DataLoader, optimizer, scheduler, and loss function<\/span><\/p>\n

dataset<\/span> <\/span>=<\/span> <\/span>PretrainingDataset<\/span>(<\/span>dataset<\/span>,<\/span> <\/span>tokenizer<\/span>,<\/span> <\/span>seq_length<\/span>,<\/span> <\/span>device<\/span>)<\/span><\/p>\n

dataloader<\/span> <\/span>=<\/span> <\/span>torch<\/span>.<\/span>utils<\/span>.<\/span>data<\/span>.<\/span>DataLoader<\/span>(<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>dataset<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>batch_size<\/span>=<\/span>batch_size<\/span>,<\/span><\/p>\n

)<\/span><\/p>\n

num_training_steps<\/span> <\/span>=<\/span> <\/span>len<\/span>(<\/span>dataloader<\/span>)<\/span> <\/span>*<\/span> <\/span>epochs<\/span><\/p>\n

print<\/span>(<\/span>f<\/span>“Number of training steps: {num_training_steps} = {len(dataloader)} * {epochs}”<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

optimizer<\/span> <\/span>=<\/span> <\/span>torch<\/span>.<\/span>optim<\/span>.<\/span>AdamW<\/span>(<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>model<\/span>.<\/span>parameters<\/span>(<\/span>)<\/span>,<\/span> <\/span>lr<\/span>=<\/span>learning_rate<\/span>,<\/span> <\/span>betas<\/span>=<\/span>(<\/span>0.9<\/span>,<\/span> <\/span>0.99<\/span>)<\/span>,<\/span> <\/span>eps<\/span>=<\/span>1e<\/span>–<\/span>8<\/span>,<\/span> <\/span>weight_decay<\/span>=<\/span>0.1<\/span>,<\/span><\/p>\n

)<\/span><\/p>\n

warmup_scheduler<\/span> <\/span>=<\/span> <\/span>lr_scheduler<\/span>.<\/span>LinearLR<\/span>(<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>optimizer<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>start_factor<\/span>=<\/span>0.1<\/span>,<\/span> <\/span>end_factor<\/span>=<\/span>1.0<\/span>,<\/span> <\/span>total_iters<\/span>=<\/span>num_warmup_steps<\/span>,<\/span><\/p>\n

)<\/span><\/p>\n

cosine_scheduler<\/span> <\/span>=<\/span> <\/span>lr_scheduler<\/span>.<\/span>CosineAnnealingLR<\/span>(<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>optimizer<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>T_max<\/span>=<\/span>num_training_steps<\/span> <\/span>–<\/span> <\/span>num_warmup_steps<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>eta_min<\/span>=<\/span>0<\/span>,<\/span><\/p>\n

)<\/span><\/p>\n

scheduler<\/span> <\/span>=<\/span> <\/span>lr_scheduler<\/span>.<\/span>SequentialLR<\/span>(<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>optimizer<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>schedulers<\/span>=<\/span>[<\/span>warmup_scheduler<\/span>,<\/span> <\/span>cosine_scheduler<\/span>]<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>milestones<\/span>=<\/span>[<\/span>num_warmup_steps<\/span>]<\/span>,<\/span><\/p>\n

)<\/span><\/p>\n

\u00a0<\/p>\n

# if checkpoint-dist dir exists, load the checkpoint to model and optimizer<\/span><\/p>\n

# Note: You should implement how to reset the epoch and step to allow correct resume<\/span><\/p>\n

if<\/span> <\/span>os<\/span>.<\/span>path<\/span>.<\/span>exists<\/span>(<\/span>“checkpoint-dist”<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>load_checkpoint<\/span>(<\/span>model<\/span>,<\/span> <\/span>optimizer<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

# Create pipeline schedule<\/span><\/p>\n

def <\/span>loss_fn<\/span>(<\/span>logits<\/span>:<\/span> <\/span>Tensor<\/span>,<\/span> <\/span>target_ids<\/span>:<\/span> <\/span>Tensor<\/span>)<\/span> <\/span>-><\/span> <\/span>Tensor<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>logits<\/span> <\/span>=<\/span> <\/span>logits<\/span>.<\/span>view<\/span>(<\/span>–<\/span>1<\/span>,<\/span> <\/span>logits<\/span>.<\/span>size<\/span>(<\/span>–<\/span>1<\/span>)<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>target_ids<\/span> <\/span>=<\/span> <\/span>target_ids<\/span>.<\/span>view<\/span>(<\/span>–<\/span>1<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>return<\/span> <\/span>F<\/span>.<\/span>cross_entropy<\/span>(<\/span>logits<\/span>,<\/span> <\/span>target_ids<\/span>,<\/span> <\/span>ignore_index<\/span>=<\/span>PAD_TOKEN_ID<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

n_microbatches<\/span> <\/span>=<\/span> <\/span>4<\/span>\u00a0\u00a0<\/span># num split per batch<\/span><\/p>\n

schedule<\/span> <\/span>=<\/span> <\/span>ScheduleGPipe<\/span>(<\/span>stage<\/span>,<\/span> <\/span>n_microbatches<\/span>=<\/span>n_microbatches<\/span>,<\/span> <\/span>loss_fn<\/span>=<\/span>loss_fn<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

# start training<\/span><\/p>\n

for<\/span> <\/span>epoch <\/span>in<\/span> <\/span>range<\/span>(<\/span>epochs<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>pbar<\/span> <\/span>=<\/span> <\/span>tqdm<\/span>.<\/span>tqdm<\/span>(<\/span>dataloader<\/span>,<\/span> <\/span>desc<\/span>=<\/span>f<\/span>“Epoch {epoch+1}\/{epochs}”<\/span>,<\/span> <\/span>disable<\/span>=<\/span>(<\/span>rank<\/span> <\/span>!=<\/span> <\/span>world_size<\/span> <\/span>–<\/span> <\/span>1<\/span>)<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>for<\/span> <\/span>batch_id<\/span>,<\/span> <\/span>batch <\/span>in<\/span> <\/span>enumerate<\/span>(<\/span>pbar<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>if<\/span> <\/span>batch_id<\/span> <\/span>%<\/span> <\/span>1000<\/span> <\/span>==<\/span> <\/span>0<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>save_checkpoint<\/span>(<\/span>model<\/span>,<\/span> <\/span>optimizer<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># zero grad before forward pass, since no explicit backward pass is called<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>optimizer<\/span>.<\/span>zero_grad<\/span>(<\/span>set_to_none<\/span>=<\/span>True<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># get batched data<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>input_ids<\/span>,<\/span> <\/span>target_ids<\/span> <\/span>=<\/span> <\/span>batch<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>if<\/span> <\/span>rank<\/span> <\/span>==<\/span> <\/span>0<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>schedule<\/span>.<\/span>step<\/span>(<\/span>input_ids<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>elif <\/span>rank<\/span> <\/span>==<\/span> <\/span>world_size<\/span> <\/span>–<\/span> <\/span>1<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>losses<\/span> <\/span>=<\/span> <\/span>[<\/span>]<\/span>\u00a0\u00a0<\/span># expects one lost per microbatch<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>logits<\/span> <\/span>=<\/span> <\/span>schedule<\/span>.<\/span>step<\/span>(<\/span>target<\/span>=<\/span>target_ids<\/span>,<\/span> <\/span>losses<\/span>=<\/span>losses<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>with <\/span>torch<\/span>.<\/span>no_grad<\/span>(<\/span>)<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>pbar<\/span>.<\/span>set_postfix<\/span>(<\/span>loss<\/span>=<\/span>sum<\/span>(<\/span>losses<\/span>)<\/span>.<\/span>item<\/span>(<\/span>)<\/span> <\/span>\/<\/span> <\/span>len<\/span>(<\/span>losses<\/span>)<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>else<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>schedule<\/span>.<\/span>step<\/span>(<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>torch<\/span>.<\/span>nn<\/span>.<\/span>utils<\/span>.<\/span>clip_grad_norm_<\/span>(<\/span>model<\/span>.<\/span>parameters<\/span>(<\/span>)<\/span>,<\/span> <\/span>1.0<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>optimizer<\/span>.<\/span>step<\/span>(<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>scheduler<\/span>.<\/span>step<\/span>(<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>pbar<\/span>.<\/span>update<\/span>(<\/span>1<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>pbar<\/span>.<\/span>close<\/span>(<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

# Save the model<\/span><\/p>\n

save_checkpoint<\/span>(<\/span>model<\/span>,<\/span> <\/span>optimizer<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

# Clean up the distributed environment<\/span><\/p>\n

dist<\/span>.<\/span>destroy_process_group<\/span>(<\/span>)<\/span><\/p>\n<\/div>\n","protected":false},"excerpt":{"rendered":"

import dataclasses import os \u00a0 import datasets import tokenizers import torch import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F import torch.optim.lr_scheduler as lr_scheduler import tqdm from torch import Tensor from torch.distributed.checkpoint import load, save from torch.distributed.checkpoint.state_dict import StateDictOptions, get_state_dict, set_state_dict from torch.distributed.pipelining import PipelineStage, ScheduleGPipe \u00a0 \u00a0 # Build the model<\/p>\n","protected":false},"author":1,"featured_media":177749,"comment_status":"open","ping_status":"open","sticky":false,"template":"","format":"standard","meta":{"footnotes":""},"categories":[42],"tags":[],"class_list":{"0":"post-177748","1":"post","2":"type-post","3":"status-publish","4":"format-standard","5":"has-post-thumbnail","7":"category-ai"},"yoast_head":"\nTrain Your Large Model on Multiple GPUs with Pipeline Parallelism - Ktromedia<\/title>\n<meta name=\"robots\" content=\"index, follow, max-snippet:-1, max-image-preview:large, max-video-preview:-1\" \/>\n<link rel=\"canonical\" href=\"http:\/\/ktromedia.com\/?p=177748\" \/>\n<meta property=\"og:locale\" content=\"en_US\" \/>\n<meta property=\"og:type\" content=\"article\" \/>\n<meta property=\"og:title\" content=\"Train Your Large Model on Multiple GPUs with Pipeline Parallelism - Ktromedia\" \/>\n<meta property=\"og:description\" content=\"import dataclasses import os \u00a0 import datasets import tokenizers import torch import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F import torch.optim.lr_scheduler as lr_scheduler import tqdm from torch import Tensor from torch.distributed.checkpoint import load, save from torch.distributed.checkpoint.state_dict import StateDictOptions, get_state_dict, set_state_dict from torch.distributed.pipelining import PipelineStage, ScheduleGPipe \u00a0 \u00a0 # Build the model\" \/>\n<meta property=\"og:url\" content=\"http:\/\/ktromedia.com\/?p=177748\" \/>\n<meta property=\"og:site_name\" content=\"Ktromedia\" \/>\n<meta property=\"article:publisher\" content=\"https:\/\/www.facebook.com\/KTROMedia\/\" \/>\n<meta property=\"article:published_time\" content=\"2026-01-04T14:18:11+00:00\" \/>\n<meta property=\"og:image\" content=\"http:\/\/ktromedia.com\/wp-content\/uploads\/2026\/01\/Train-Your-Large-Model-on-Multiple-GPUs-with-Pipeline-Parallelism.jpg\" \/>\n\t<meta property=\"og:image:width\" content=\"2560\" \/>\n\t<meta property=\"og:image:height\" content=\"1440\" \/>\n\t<meta property=\"og:image:type\" content=\"image\/jpeg\" \/>\n<meta name=\"author\" content=\"KTRO TEAM\" \/>\n<meta name=\"twitter:card\" content=\"summary_large_image\" \/>\n<meta name=\"twitter:label1\" content=\"Written by\" \/>\n\t<meta name=\"twitter:data1\" content=\"KTRO TEAM\" \/>\n\t<meta name=\"twitter:label2\" content=\"Est. reading time\" \/>\n\t<meta name=\"twitter:data2\" content=\"11 minutes\" \/>\n<script type=\"application\/ld+json\" class=\"yoast-schema-graph\">{\"@context\":\"https:\/\/schema.org\",\"@graph\":[{\"@type\":\"Article\",\"@id\":\"http:\/\/ktromedia.com\/?p=177748#article\",\"isPartOf\":{\"@id\":\"http:\/\/ktromedia.com\/?p=177748\"},\"author\":{\"name\":\"KTRO TEAM\",\"@id\":\"https:\/\/ktromedia.com\/#\/schema\/person\/612bf2fbac107722ea365932cdd35f5b\"},\"headline\":\"Train Your Large Model on Multiple GPUs with Pipeline Parallelism\",\"datePublished\":\"2026-01-04T14:18:11+00:00\",\"mainEntityOfPage\":{\"@id\":\"http:\/\/ktromedia.com\/?p=177748\"},\"wordCount\":2145,\"commentCount\":0,\"publisher\":{\"@id\":\"https:\/\/ktromedia.com\/#organization\"},\"image\":{\"@id\":\"http:\/\/ktromedia.com\/?p=177748#primaryimage\"},\"thumbnailUrl\":\"http:\/\/ktromedia.com\/wp-content\/uploads\/2026\/01\/Train-Your-Large-Model-on-Multiple-GPUs-with-Pipeline-Parallelism.jpg\",\"articleSection\":[\"\u4eba\u5de5\u667a\u80fd\"],\"inLanguage\":\"en-US\",\"potentialAction\":[{\"@type\":\"CommentAction\",\"name\":\"Comment\",\"target\":[\"http:\/\/ktromedia.com\/?p=177748#respond\"]}]},{\"@type\":\"WebPage\",\"@id\":\"http:\/\/ktromedia.com\/?p=177748\",\"url\":\"http:\/\/ktromedia.com\/?p=177748\",\"name\":\"Train Your Large Model on Multiple GPUs with Pipeline Parallelism - Ktromedia\",\"isPartOf\":{\"@id\":\"https:\/\/ktromedia.com\/#website\"},\"primaryImageOfPage\":{\"@id\":\"http:\/\/ktromedia.com\/?p=177748#primaryimage\"},\"image\":{\"@id\":\"http:\/\/ktromedia.com\/?p=177748#primaryimage\"},\"thumbnailUrl\":\"http:\/\/ktromedia.com\/wp-content\/uploads\/2026\/01\/Train-Your-Large-Model-on-Multiple-GPUs-with-Pipeline-Parallelism.jpg\",\"datePublished\":\"2026-01-04T14:18:11+00:00\",\"breadcrumb\":{\"@id\":\"http:\/\/ktromedia.com\/?p=177748#breadcrumb\"},\"inLanguage\":\"en-US\",\"potentialAction\":[{\"@type\":\"ReadAction\",\"target\":[\"http:\/\/ktromedia.com\/?p=177748\"]}]},{\"@type\":\"ImageObject\",\"inLanguage\":\"en-US\",\"@id\":\"http:\/\/ktromedia.com\/?p=177748#primaryimage\",\"url\":\"http:\/\/ktromedia.com\/wp-content\/uploads\/2026\/01\/Train-Your-Large-Model-on-Multiple-GPUs-with-Pipeline-Parallelism.jpg\",\"contentUrl\":\"http:\/\/ktromedia.com\/wp-content\/uploads\/2026\/01\/Train-Your-Large-Model-on-Multiple-GPUs-with-Pipeline-Parallelism.jpg\",\"width\":2560,\"height\":1440},{\"@type\":\"BreadcrumbList\",\"@id\":\"http:\/\/ktromedia.com\/?p=177748#breadcrumb\",\"itemListElement\":[{\"@type\":\"ListItem\",\"position\":1,\"name\":\"Home\",\"item\":\"https:\/\/ktromedia.com\/\"},{\"@type\":\"ListItem\",\"position\":2,\"name\":\"Train Your Large Model on Multiple GPUs with Pipeline Parallelism\"}]},{\"@type\":\"WebSite\",\"@id\":\"https:\/\/ktromedia.com\/#website\",\"url\":\"https:\/\/ktromedia.com\/\",\"name\":\"Ktromedia\",\"description\":\"KTRO MEDIA Crypto News\",\"publisher\":{\"@id\":\"https:\/\/ktromedia.com\/#organization\"},\"potentialAction\":[{\"@type\":\"SearchAction\",\"target\":{\"@type\":\"EntryPoint\",\"urlTemplate\":\"https:\/\/ktromedia.com\/?s={search_term_string}\"},\"query-input\":{\"@type\":\"PropertyValueSpecification\",\"valueRequired\":true,\"valueName\":\"search_term_string\"}}],\"inLanguage\":\"en-US\"},{\"@type\":\"Organization\",\"@id\":\"https:\/\/ktromedia.com\/#organization\",\"name\":\"Ktromedia\",\"url\":\"https:\/\/ktromedia.com\/\",\"logo\":{\"@type\":\"ImageObject\",\"inLanguage\":\"en-US\",\"@id\":\"https:\/\/ktromedia.com\/#\/schema\/logo\/image\/\",\"url\":\"https:\/\/ktromedia.com\/wp-content\/uploads\/2025\/11\/ktroicon.png\",\"contentUrl\":\"https:\/\/ktromedia.com\/wp-content\/uploads\/2025\/11\/ktroicon.png\",\"width\":250,\"height\":250,\"caption\":\"Ktromedia\"},\"image\":{\"@id\":\"https:\/\/ktromedia.com\/#\/schema\/logo\/image\/\"},\"sameAs\":[\"https:\/\/www.facebook.com\/KTROMedia\/\",\"https:\/\/www.linkedin.com\/company\/ktro-media\/\",\"https:\/\/t.me\/ktrogroup\"]},{\"@type\":\"Person\",\"@id\":\"https:\/\/ktromedia.com\/#\/schema\/person\/612bf2fbac107722ea365932cdd35f5b\",\"name\":\"KTRO TEAM\",\"image\":{\"@type\":\"ImageObject\",\"inLanguage\":\"en-US\",\"@id\":\"https:\/\/ktromedia.com\/#\/schema\/person\/image\/\",\"url\":\"http:\/\/ktromedia.com\/wp-content\/uploads\/2025\/10\/cropped-Untitled-design-7-1-150x150.png\",\"contentUrl\":\"http:\/\/ktromedia.com\/wp-content\/uploads\/2025\/10\/cropped-Untitled-design-7-1-150x150.png\",\"caption\":\"KTRO TEAM\"},\"description\":\"KTRO MEDIA \u662f\u4e00\u5bb6\u5168\u7403\u6027\u7684\u534e\u6587WEB3\u5a92\u4f53\u516c\u53f8\u3002\u6211\u4eec\u81f4\u529b\u4e8e\u4e3a\u533a\u5757\u94fe\u548c\u91d1\u878d\u79d1\u6280\u9886\u57df\u63d0\u4f9b\u6700\u65b0\u7684\u65b0\u95fb\u3001\u89c1\u89e3\u548c\u8d8b\u52bf\u5206\u6790\u3002\u6211\u4eec\u7684\u5b97\u65e8\u662f\u4e3a\u5168\u7403\u7528\u6237\u63d0\u4f9b\u9ad8\u8d28\u91cf\u3001\u5168\u9762\u7684\u8d44\u8baf\u670d\u52a1\uff0c\u8ba9\u4ed6\u4eec\u66f4\u597d\u5730\u4e86\u89e3\u533a\u5757\u94fe\u548c\u91d1\u878d\u79d1\u6280\u884c\u4e1a\u7684\u6700\u65b0\u52a8\u6001\u3002\u6211\u4eec\u4e5f\u5e0c\u671b\u80fd\u5e2e\u5230\u66f4\u591a\u4f18\u79c0\u7684WEB3\u4ea7\u54c1\u627e\u5230\u66f4\u591a\u66f4\u597d\u7684\u8d44\u6e90\u597d\u8ba9\u8fd9\u9886\u57df\u53d8\u5f97\u66f4\u6210\u719f\u3002 \u6211\u4eec\u7684\u62a5\u9053\u8303\u56f4\u6db5\u76d6\u4e86\u533a\u5757\u94fe\u3001\u52a0\u5bc6\u8d27\u5e01\u3001\u667a\u80fd\u5408\u7ea6\u3001DeFi\u3001NFT \u548c Web3 \u751f\u6001\u7cfb\u7edf\u7b49\u9886\u57df\u3002\u6211\u4eec\u7684\u62a5\u9053\u4e0d\u4ec5\u6765\u81ea\u884c\u4e1a\u5185\u7684\u4e13\u5bb6\uff0c\u5148\u950b\u8005\u4e5f\u5305\u62ec\u4e86\u6211\u4eec\u81ea\u5df1\u7684\u5206\u6790\u548c\u89c2\u70b9\u3002\u6211\u4eec\u5728\u5404\u4e2a\u56fd\u5bb6\u548c\u5730\u533a\u90fd\u8bbe\u6709\u56e2\u961f\uff0c\u4e3a\u8bfb\u8005\u63d0\u4f9b\u672c\u5730\u5316\u7684\u62a5\u9053\u548c\u5206\u6790\u3002 \u9664\u4e86\u65b0\u95fb\u62a5\u9053\uff0c\u6211\u4eec\u8fd8\u63d0\u4f9b\u5e02\u573a\u7814\u7a76\u548c\u54a8\u8be2\u670d\u52a1\u3002\u6211\u4eec\u7684\u4e13\u4e1a\u56e2\u961f\u53ef\u4ee5\u4e3a\u60a8\u63d0\u4f9b\u6709\u5173\u533a\u5757\u94fe\u548c\u91d1\u878d\u79d1\u6280\u884c\u4e1a\u7684\u6df1\u5165\u5206\u6790\u548c\u5e02\u573a\u8d8b\u52bf\uff0c\u5e2e\u52a9\u60a8\u505a\u51fa\u66f4\u660e\u667a\u7684\u6295\u8d44\u51b3\u7b56\u3002 \u6211\u4eec\u7684\u4f7f\u547d\u662f\u6210\u4e3a\u5168\u7403\u534e\u6587\u533a\u5757\u94fe\u548c\u91d1\u878d\u79d1\u6280\u884c\u4e1a\u6700\u53d7\u4fe1\u8d56\u7684\u4fe1\u606f\u6765\u6e90\u4e4b\u4e00\u3002\u6211\u4eec\u5c06\u7ee7\u7eed\u4e0d\u65ad\u52aa\u529b\uff0c\u4e3a\u8bfb\u8005\u63d0\u4f9b\u6700\u65b0\u3001\u6700\u5168\u9762\u3001\u6700\u53ef\u9760\u7684\u4fe1\u606f\u670d\u52a1\u3002\",\"sameAs\":[\"https:\/\/ktromedia.com\"],\"url\":\"http:\/\/ktromedia.com\/?author=1\"}]}<\/script>\n<!-- \/ Yoast SEO plugin. -->","yoast_head_json":{"title":"Train Your Large Model on Multiple GPUs with Pipeline Parallelism - Ktromedia","robots":{"index":"index","follow":"follow","max-snippet":"max-snippet:-1","max-image-preview":"max-image-preview:large","max-video-preview":"max-video-preview:-1"},"canonical":"http:\/\/ktromedia.com\/?p=177748","og_locale":"en_US","og_type":"article","og_title":"Train Your Large Model on Multiple GPUs with Pipeline Parallelism - Ktromedia","og_description":"import dataclasses import os \u00a0 import datasets import tokenizers import torch import torch.distributed as dist import torch.nn as nn import torch.nn.functional as F import torch.optim.lr_scheduler as lr_scheduler import tqdm from torch import Tensor from torch.distributed.checkpoint import load, save from torch.distributed.checkpoint.state_dict import StateDictOptions, get_state_dict, set_state_dict from torch.distributed.pipelining import PipelineStage, ScheduleGPipe \u00a0 \u00a0 # Build the model","og_url":"http:\/\/ktromedia.com\/?p=177748","og_site_name":"Ktromedia","article_publisher":"https:\/\/www.facebook.com\/KTROMedia\/","article_published_time":"2026-01-04T14:18:11+00:00","og_image":[{"width":2560,"height":1440,"url":"http:\/\/ktromedia.com\/wp-content\/uploads\/2026\/01\/Train-Your-Large-Model-on-Multiple-GPUs-with-Pipeline-Parallelism.jpg","type":"image\/jpeg"}],"author":"KTRO TEAM","twitter_card":"summary_large_image","twitter_misc":{"Written by":"KTRO TEAM","Est. reading time":"11 minutes"},"schema":{"@context":"https:\/\/schema.org","@graph":[{"@type":"Article","@id":"http:\/\/ktromedia.com\/?p=177748#article","isPartOf":{"@id":"http:\/\/ktromedia.com\/?p=177748"},"author":{"name":"KTRO TEAM","@id":"https:\/\/ktromedia.com\/#\/schema\/person\/612bf2fbac107722ea365932cdd35f5b"},"headline":"Train Your Large Model on Multiple GPUs with Pipeline Parallelism","datePublished":"2026-01-04T14:18:11+00:00","mainEntityOfPage":{"@id":"http:\/\/ktromedia.com\/?p=177748"},"wordCount":2145,"commentCount":0,"publisher":{"@id":"https:\/\/ktromedia.com\/#organization"},"image":{"@id":"http:\/\/ktromedia.com\/?p=177748#primaryimage"},"thumbnailUrl":"http:\/\/ktromedia.com\/wp-content\/uploads\/2026\/01\/Train-Your-Large-Model-on-Multiple-GPUs-with-Pipeline-Parallelism.jpg","articleSection":["\u4eba\u5de5\u667a\u80fd"],"inLanguage":"en-US","potentialAction":[{"@type":"CommentAction","name":"Comment","target":["http:\/\/ktromedia.com\/?p=177748#respond"]}]},{"@type":"WebPage","@id":"http:\/\/ktromedia.com\/?p=177748","url":"http:\/\/ktromedia.com\/?p=177748","name":"Train Your Large Model on Multiple GPUs with Pipeline Parallelism - Ktromedia","isPartOf":{"@id":"https:\/\/ktromedia.com\/#website"},"primaryImageOfPage":{"@id":"http:\/\/ktromedia.com\/?p=177748#primaryimage"},"image":{"@id":"http:\/\/ktromedia.com\/?p=177748#primaryimage"},"thumbnailUrl":"http:\/\/ktromedia.com\/wp-content\/uploads\/2026\/01\/Train-Your-Large-Model-on-Multiple-GPUs-with-Pipeline-Parallelism.jpg","datePublished":"2026-01-04T14:18:11+00:00","breadcrumb":{"@id":"http:\/\/ktromedia.com\/?p=177748#breadcrumb"},"inLanguage":"en-US","potentialAction":[{"@type":"ReadAction","target":["http:\/\/ktromedia.com\/?p=177748"]}]},{"@type":"ImageObject","inLanguage":"en-US","@id":"http:\/\/ktromedia.com\/?p=177748#primaryimage","url":"http:\/\/ktromedia.com\/wp-content\/uploads\/2026\/01\/Train-Your-Large-Model-on-Multiple-GPUs-with-Pipeline-Parallelism.jpg","contentUrl":"http:\/\/ktromedia.com\/wp-content\/uploads\/2026\/01\/Train-Your-Large-Model-on-Multiple-GPUs-with-Pipeline-Parallelism.jpg","width":2560,"height":1440},{"@type":"BreadcrumbList","@id":"http:\/\/ktromedia.com\/?p=177748#breadcrumb","itemListElement":[{"@type":"ListItem","position":1,"name":"Home","item":"https:\/\/ktromedia.com\/"},{"@type":"ListItem","position":2,"name":"Train Your Large Model on Multiple GPUs with Pipeline Parallelism"}]},{"@type":"WebSite","@id":"https:\/\/ktromedia.com\/#website","url":"https:\/\/ktromedia.com\/","name":"Ktromedia","description":"KTRO MEDIA Crypto News","publisher":{"@id":"https:\/\/ktromedia.com\/#organization"},"potentialAction":[{"@type":"SearchAction","target":{"@type":"EntryPoint","urlTemplate":"https:\/\/ktromedia.com\/?s={search_term_string}"},"query-input":{"@type":"PropertyValueSpecification","valueRequired":true,"valueName":"search_term_string"}}],"inLanguage":"en-US"},{"@type":"Organization","@id":"https:\/\/ktromedia.com\/#organization","name":"Ktromedia","url":"https:\/\/ktromedia.com\/","logo":{"@type":"ImageObject","inLanguage":"en-US","@id":"https:\/\/ktromedia.com\/#\/schema\/logo\/image\/","url":"https:\/\/ktromedia.com\/wp-content\/uploads\/2025\/11\/ktroicon.png","contentUrl":"https:\/\/ktromedia.com\/wp-content\/uploads\/2025\/11\/ktroicon.png","width":250,"height":250,"caption":"Ktromedia"},"image":{"@id":"https:\/\/ktromedia.com\/#\/schema\/logo\/image\/"},"sameAs":["https:\/\/www.facebook.com\/KTROMedia\/","https:\/\/www.linkedin.com\/company\/ktro-media\/","https:\/\/t.me\/ktrogroup"]},{"@type":"Person","@id":"https:\/\/ktromedia.com\/#\/schema\/person\/612bf2fbac107722ea365932cdd35f5b","name":"KTRO TEAM","image":{"@type":"ImageObject","inLanguage":"en-US","@id":"https:\/\/ktromedia.com\/#\/schema\/person\/image\/","url":"http:\/\/ktromedia.com\/wp-content\/uploads\/2025\/10\/cropped-Untitled-design-7-1-150x150.png","contentUrl":"http:\/\/ktromedia.com\/wp-content\/uploads\/2025\/10\/cropped-Untitled-design-7-1-150x150.png","caption":"KTRO TEAM"},"description":"KTRO MEDIA \u662f\u4e00\u5bb6\u5168\u7403\u6027\u7684\u534e\u6587WEB3\u5a92\u4f53\u516c\u53f8\u3002\u6211\u4eec\u81f4\u529b\u4e8e\u4e3a\u533a\u5757\u94fe\u548c\u91d1\u878d\u79d1\u6280\u9886\u57df\u63d0\u4f9b\u6700\u65b0\u7684\u65b0\u95fb\u3001\u89c1\u89e3\u548c\u8d8b\u52bf\u5206\u6790\u3002\u6211\u4eec\u7684\u5b97\u65e8\u662f\u4e3a\u5168\u7403\u7528\u6237\u63d0\u4f9b\u9ad8\u8d28\u91cf\u3001\u5168\u9762\u7684\u8d44\u8baf\u670d\u52a1\uff0c\u8ba9\u4ed6\u4eec\u66f4\u597d\u5730\u4e86\u89e3\u533a\u5757\u94fe\u548c\u91d1\u878d\u79d1\u6280\u884c\u4e1a\u7684\u6700\u65b0\u52a8\u6001\u3002\u6211\u4eec\u4e5f\u5e0c\u671b\u80fd\u5e2e\u5230\u66f4\u591a\u4f18\u79c0\u7684WEB3\u4ea7\u54c1\u627e\u5230\u66f4\u591a\u66f4\u597d\u7684\u8d44\u6e90\u597d\u8ba9\u8fd9\u9886\u57df\u53d8\u5f97\u66f4\u6210\u719f\u3002 \u6211\u4eec\u7684\u62a5\u9053\u8303\u56f4\u6db5\u76d6\u4e86\u533a\u5757\u94fe\u3001\u52a0\u5bc6\u8d27\u5e01\u3001\u667a\u80fd\u5408\u7ea6\u3001DeFi\u3001NFT \u548c Web3 \u751f\u6001\u7cfb\u7edf\u7b49\u9886\u57df\u3002\u6211\u4eec\u7684\u62a5\u9053\u4e0d\u4ec5\u6765\u81ea\u884c\u4e1a\u5185\u7684\u4e13\u5bb6\uff0c\u5148\u950b\u8005\u4e5f\u5305\u62ec\u4e86\u6211\u4eec\u81ea\u5df1\u7684\u5206\u6790\u548c\u89c2\u70b9\u3002\u6211\u4eec\u5728\u5404\u4e2a\u56fd\u5bb6\u548c\u5730\u533a\u90fd\u8bbe\u6709\u56e2\u961f\uff0c\u4e3a\u8bfb\u8005\u63d0\u4f9b\u672c\u5730\u5316\u7684\u62a5\u9053\u548c\u5206\u6790\u3002 \u9664\u4e86\u65b0\u95fb\u62a5\u9053\uff0c\u6211\u4eec\u8fd8\u63d0\u4f9b\u5e02\u573a\u7814\u7a76\u548c\u54a8\u8be2\u670d\u52a1\u3002\u6211\u4eec\u7684\u4e13\u4e1a\u56e2\u961f\u53ef\u4ee5\u4e3a\u60a8\u63d0\u4f9b\u6709\u5173\u533a\u5757\u94fe\u548c\u91d1\u878d\u79d1\u6280\u884c\u4e1a\u7684\u6df1\u5165\u5206\u6790\u548c\u5e02\u573a\u8d8b\u52bf\uff0c\u5e2e\u52a9\u60a8\u505a\u51fa\u66f4\u660e\u667a\u7684\u6295\u8d44\u51b3\u7b56\u3002 \u6211\u4eec\u7684\u4f7f\u547d\u662f\u6210\u4e3a\u5168\u7403\u534e\u6587\u533a\u5757\u94fe\u548c\u91d1\u878d\u79d1\u6280\u884c\u4e1a\u6700\u53d7\u4fe1\u8d56\u7684\u4fe1\u606f\u6765\u6e90\u4e4b\u4e00\u3002\u6211\u4eec\u5c06\u7ee7\u7eed\u4e0d\u65ad\u52aa\u529b\uff0c\u4e3a\u8bfb\u8005\u63d0\u4f9b\u6700\u65b0\u3001\u6700\u5168\u9762\u3001\u6700\u53ef\u9760\u7684\u4fe1\u606f\u670d\u52a1\u3002","sameAs":["https:\/\/ktromedia.com"],"url":"http:\/\/ktromedia.com\/?author=1"}]}},"_links":{"self":[{"href":"http:\/\/ktromedia.com\/index.php?rest_route=\/wp\/v2\/posts\/177748","targetHints":{"allow":["GET"]}}],"collection":[{"href":"http:\/\/ktromedia.com\/index.php?rest_route=\/wp\/v2\/posts"}],"about":[{"href":"http:\/\/ktromedia.com\/index.php?rest_route=\/wp\/v2\/types\/post"}],"author":[{"embeddable":true,"href":"http:\/\/ktromedia.com\/index.php?rest_route=\/wp\/v2\/users\/1"}],"replies":[{"embeddable":true,"href":"http:\/\/ktromedia.com\/index.php?rest_route=%2Fwp%2Fv2%2Fcomments&post=177748"}],"version-history":[{"count":1,"href":"http:\/\/ktromedia.com\/index.php?rest_route=\/wp\/v2\/posts\/177748\/revisions"}],"predecessor-version":[{"id":177750,"href":"http:\/\/ktromedia.com\/index.php?rest_route=\/wp\/v2\/posts\/177748\/revisions\/177750"}],"wp:featuredmedia":[{"embeddable":true,"href":"http:\/\/ktromedia.com\/index.php?rest_route=\/wp\/v2\/media\/177749"}],"wp:attachment":[{"href":"http:\/\/ktromedia.com\/index.php?rest_route=%2Fwp%2Fv2%2Fmedia&parent=177748"}],"wp:term":[{"taxonomy":"category","embeddable":true,"href":"http:\/\/ktromedia.com\/index.php?rest_route=%2Fwp%2Fv2%2Fcategories&post=177748"},{"taxonomy":"post_tag","embeddable":true,"href":"http:\/\/ktromedia.com\/index.php?rest_route=%2Fwp%2Fv2%2Ftags&post=177748"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}