{"id":177895,"date":"2026-01-12T14:27:19","date_gmt":"2026-01-12T14:27:19","guid":{"rendered":"https:\/\/ktromedia.com\/?p=177895"},"modified":"2026-01-12T14:27:19","modified_gmt":"2026-01-12T14:27:19","slug":"training-a-model-on-multiple-gpus-with-data-parallelism","status":"publish","type":"post","link":"http:\/\/ktromedia.com\/?p=177895","title":{"rendered":"Training a Model on Multiple GPUs with Data 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>tqdm<\/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

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

from <\/span>torch<\/span>.<\/span>nn<\/span>.<\/span>parallel <\/span>import <\/span>DistributedDataParallel <\/span>as<\/span> <\/span>DDP<\/span><\/p>\n

from <\/span>torch<\/span>.<\/span>utils<\/span>.<\/span>data<\/span>.<\/span>distributed <\/span>import <\/span>DistributedSampler<\/span><\/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>attn_mask<\/span>:<\/span> <\/span>Tensor<\/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>attn_mask<\/span>=<\/span>attn_mask<\/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>attn_mask<\/span>:<\/span> <\/span>Tensor<\/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> <\/span>attn_mask<\/span>=<\/span>attn_mask<\/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>rotary_emb<\/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>ModuleList<\/span>(<\/span>[<\/span>LlamaDecoderLayer<\/span>(<\/span>config<\/span>)<\/span> <\/span>for<\/span> <\/span>_<\/span> <\/span>in<\/span> <\/span>range<\/span>(<\/span>config<\/span>.<\/span>num_hidden_layers<\/span>)<\/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>attn_mask<\/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>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># 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>layer <\/span>in<\/span> <\/span>self<\/span>.<\/span>layers<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>hidden_states<\/span> <\/span>=<\/span> <\/span>layer<\/span>(<\/span>hidden_states<\/span>,<\/span> <\/span>rope<\/span>=<\/span>self<\/span>.<\/span>rotary_emb<\/span>,<\/span> <\/span>attn_mask<\/span>=<\/span>attn_mask<\/span>)<\/span><\/p>\n

\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<\/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>attn_mask<\/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> <\/span>attn_mask<\/span>)<\/span><\/p>\n

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

\u00a0<\/p>\n

\u00a0<\/p>\n

def <\/span>create_causal_mask<\/span>(<\/span>batch<\/span>:<\/span> <\/span>Tensor<\/span>,<\/span> <\/span>dtype<\/span>:<\/span> <\/span>torch<\/span>.<\/span>dtype<\/span> <\/span>=<\/span> <\/span>torch<\/span>.<\/span>float32<\/span>)<\/span> <\/span>-><\/span> <\/span>Tensor<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>“”<\/span>“Create a causal mask for self-attention.<\/span><\/p>\n

\u00a0<\/p>\n

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

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0batch: Batch of sequences, shape (batch_size, seq_len)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0dtype: Data type of the mask<\/span><\/p>\n

\u00a0<\/p>\n

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

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0Causal mask of shape (seq_len, seq_len)<\/span><\/p>\n

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

\u00a0\u00a0\u00a0\u00a0<\/span>batch_size<\/span>,<\/span> <\/span>seq_len<\/span> <\/span>=<\/span> <\/span>batch<\/span>.<\/span>shape<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>mask<\/span> <\/span>=<\/span> <\/span>torch<\/span>.<\/span>full<\/span>(<\/span>(<\/span>seq_len<\/span>,<\/span> <\/span>seq_len<\/span>)<\/span>,<\/span> <\/span>float<\/span>(<\/span>‘-inf’<\/span>)<\/span>,<\/span> <\/span>device<\/span>=<\/span>batch<\/span>.<\/span>device<\/span>,<\/span> <\/span>dtype<\/span>=<\/span>dtype<\/span>)<\/span> <\/span>\\<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>.<\/span>triu<\/span>(<\/span>diagonal<\/span>=<\/span>1<\/span>)<\/span><\/p>\n

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

\u00a0<\/p>\n

\u00a0<\/p>\n

def <\/span>create_padding_mask<\/span>(<\/span>batch<\/span>:<\/span> <\/span>Tensor<\/span>,<\/span> <\/span>padding_token_id<\/span>:<\/span> <\/span>int<\/span>,<\/span> <\/span>dtype<\/span>:<\/span> <\/span>torch<\/span>.<\/span>dtype<\/span> <\/span>=<\/span> <\/span>torch<\/span>.<\/span>float32<\/span>)<\/span> <\/span>-><\/span> <\/span>Tensor<\/span>:<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>“”<\/span>“Create a padding mask for a batch of sequences for self-attention.<\/span><\/p>\n

\u00a0<\/p>\n

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

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0batch: Batch of sequences, shape (batch_size, seq_len)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0padding_token_id: ID of the padding token<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0dtype: Data type of the mask<\/span><\/p>\n

\u00a0<\/p>\n

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

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0Padding mask of shape (batch_size, 1, seq_len, seq_len)<\/span><\/p>\n

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

\u00a0\u00a0\u00a0\u00a0<\/span>padded<\/span> <\/span>=<\/span> <\/span>torch<\/span>.<\/span>zeros_like<\/span>(<\/span>batch<\/span>,<\/span> <\/span>device<\/span>=<\/span>batch<\/span>.<\/span>device<\/span>,<\/span> <\/span>dtype<\/span>=<\/span>dtype<\/span>)<\/span> <\/span>\\<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>.<\/span>masked_fill<\/span>(<\/span>batch<\/span> <\/span>==<\/span> <\/span>padding_token_id<\/span>,<\/span> <\/span>float<\/span>(<\/span>‘-inf’<\/span>)<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>mask<\/span> <\/span>=<\/span> <\/span>padded<\/span>[<\/span>:<\/span>,<\/span>:<\/span>,<\/span>None<\/span>]<\/span> <\/span>+<\/span> <\/span>padded<\/span>[<\/span>:<\/span>,<\/span>None<\/span>,<\/span>:<\/span>]<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>return<\/span> <\/span>mask<\/span>[<\/span>:<\/span>,<\/span> <\/span>None<\/span>,<\/span> <\/span>:<\/span>,<\/span> <\/span>:<\/span>]<\/span><\/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><\/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>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><\/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><\/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

# Load the tokenizer<\/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

\u00a0<\/p>\n

# Load the dataset<\/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: {device}”<\/span>)<\/span><\/p>\n

\u00a0<\/p>\n

# Create pretraining model with default config, then wrap it in DDP<\/span><\/p>\n

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

model<\/span> <\/span>=<\/span> <\/span>LlamaForPretraining<\/span>(<\/span>model_config<\/span>)<\/span>.<\/span>to<\/span>(<\/span>rank<\/span>)<\/span><\/p>\n

model<\/span> <\/span>=<\/span> <\/span>DDP<\/span>(<\/span>model<\/span>,<\/span> <\/span>device_ids<\/span>=<\/span>[<\/span>local_rank<\/span>]<\/span>)<\/span><\/p>\n

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

\u00a0<\/p>\n

# print the model size<\/span><\/p>\n

print<\/span>(<\/span>f<\/span>“Model parameters size: {sum(p.numel() for p in model.parameters()) \/ 1024**2:.2f} M”<\/span>)<\/span><\/p>\n

print<\/span>(<\/span>f<\/span>“Model buffers size: {sum(p.numel() for p in model.buffers()) \/ 1024**2:.2f} M”<\/span>)<\/span><\/p>\n

print<\/span>(<\/span>f<\/span>“Model precision(s): {set([x.dtype for x in model.state_dict().values()])}”<\/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><\/p>\n

sampler<\/span> <\/span>=<\/span> <\/span>DistributedSampler<\/span>(<\/span>dataset<\/span>,<\/span> <\/span>shuffle<\/span>=<\/span>False<\/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

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

\u00a0\u00a0\u00a0\u00a0<\/span>pin_memory<\/span>=<\/span>True<\/span>,<\/span>\u00a0\u00a0<\/span># optional<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>shuffle<\/span>=<\/span>False<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>num_workers<\/span>=<\/span>world_size<\/span>,<\/span><\/p>\n

)<\/span><\/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><\/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

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<\/span>_<\/span>steps<\/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><\/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><\/p>\n

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

loss_fn<\/span> <\/span>=<\/span> <\/span>nn<\/span>.<\/span>CrossEntropyLoss<\/span>(<\/span>ignore_index<\/span>=<\/span>PAD_TOKEN_ID<\/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><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>sampler<\/span>.<\/span>set_epoch<\/span>(<\/span>epoch<\/span>)<\/span>\u00a0\u00a0 <\/span># required for shuffling only<\/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>and<\/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># checkpoint the model and optimizer state, only on rank 0 process<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>torch<\/span>.<\/span>save<\/span>(<\/span>{<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>“model”<\/span>:<\/span> <\/span>model<\/span>.<\/span>module<\/span>.<\/span>state_dict<\/span>(<\/span>)<\/span> <\/span>if<\/span> <\/span>isinstance<\/span>(<\/span>model<\/span>,<\/span> <\/span>DDP<\/span>)<\/span> <\/span>else<\/span> <\/span>model<\/span>.<\/span>state_dict<\/span>(<\/span>)<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>“optimizer”<\/span>:<\/span> <\/span>optimizer<\/span>.<\/span>state_dict<\/span>(<\/span>)<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>“scheduler”<\/span>:<\/span> <\/span>scheduler<\/span>.<\/span>state_dict<\/span>(<\/span>)<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>“epoch”<\/span>:<\/span> <\/span>epoch<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>“batch”<\/span>:<\/span> <\/span>batch_id<\/span>,<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>}<\/span>,<\/span> <\/span>f<\/span>“llama_pretraining_checkpoint.pth”<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># get batched data, move from CPU to GPU<\/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>input_ids<\/span> <\/span>=<\/span> <\/span>input_ids<\/span>.<\/span>to<\/span>(<\/span>device<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>target_ids<\/span> <\/span>=<\/span> <\/span>target_ids<\/span>.<\/span>to<\/span>(<\/span>device<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># create attention mask: causal mask + padding mask<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>attn_mask<\/span> <\/span>=<\/span> <\/span>create_causal_mask<\/span>(<\/span>input_ids<\/span>)<\/span> <\/span>+<\/span> <\/span>\\<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>create_padding_mask<\/span>(<\/span>input_ids<\/span>,<\/span> <\/span>PAD_TOKEN_ID<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># extract output from model<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>logits<\/span> <\/span>=<\/span> <\/span>model<\/span>(<\/span>input_ids<\/span>,<\/span> <\/span>attn_mask<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># compute loss: cross-entropy between logits and target, ignoring padding tokens<\/span><\/p>\n

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

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span># backward with loss and gradient clipping by L2 norm to 1.0<\/span><\/p>\n

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

\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0\u00a0<\/span>loss<\/span>.<\/span>backward<\/span>(<\/span>)<\/span><\/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>set_postfix<\/span>(<\/span>loss<\/span>=<\/span>loss<\/span>.<\/span>item<\/span>(<\/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

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

\u00a0\u00a0\u00a0\u00a0<\/span>torch<\/span>.<\/span>save<\/span>(<\/span>model<\/span>.<\/span>state_dict<\/span>(<\/span>)<\/span>,<\/span> <\/span>“llama_pretraining_model.pth”<\/span>)<\/span><\/p>\n

\u00a0\u00a0\u00a0\u00a0<\/span>torch<\/span>.<\/span>save<\/span>(<\/span>model<\/span>.<\/span>base_model<\/span>.<\/span>state_dict<\/span>(<\/span>)<\/span>,<\/span> <\/span>“llama_model.pth”<\/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 tqdm 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 from torch import Tensor from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data.distributed import DistributedSampler \u00a0 # Build the model @dataclasses.dataclass class LlamaConfig: \u00a0\u00a0\u00a0\u00a0“”“Define Llama model hyperparameters.”“”<\/p>\n","protected":false},"author":1,"featured_media":177896,"comment_status":"open","ping_status":"open","sticky":false,"template":"","format":"standard","meta":{"footnotes":""},"categories":[42],"tags":[],"class_list":{"0":"post-177895","1":"post","2":"type-post","3":"status-publish","4":"format-standard","5":"has-post-thumbnail","7":"category-ai"},"yoast_head":"\nTraining a Model on Multiple GPUs with Data 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=177895\" \/>\n<meta property=\"og:locale\" content=\"en_US\" \/>\n<meta property=\"og:type\" content=\"article\" \/>\n<meta property=\"og:title\" content=\"Training a Model on Multiple GPUs with Data Parallelism - Ktromedia\" \/>\n<meta property=\"og:description\" content=\"import dataclasses import os \u00a0 import datasets import tqdm 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 from torch import Tensor from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data.distributed import DistributedSampler \u00a0 # Build the model @dataclasses.dataclass class LlamaConfig: \u00a0\u00a0\u00a0\u00a0“”“Define Llama model hyperparameters.”“”\" \/>\n<meta property=\"og:url\" content=\"http:\/\/ktromedia.com\/?p=177895\" \/>\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-12T14:27:19+00:00\" \/>\n<meta property=\"og:image\" content=\"http:\/\/ktromedia.com\/wp-content\/uploads\/2026\/01\/Training-a-Model-on-Multiple-GPUs-with-Data-Parallelism.jpg\" \/>\n\t<meta property=\"og:image:width\" content=\"2560\" \/>\n\t<meta property=\"og:image:height\" content=\"1656\" \/>\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=\"10 minutes\" \/>\n<script type=\"application\/ld+json\" class=\"yoast-schema-graph\">{\"@context\":\"https:\/\/schema.org\",\"@graph\":[{\"@type\":\"Article\",\"@id\":\"http:\/\/ktromedia.com\/?p=177895#article\",\"isPartOf\":{\"@id\":\"http:\/\/ktromedia.com\/?p=177895\"},\"author\":{\"name\":\"KTRO TEAM\",\"@id\":\"https:\/\/ktromedia.com\/#\/schema\/person\/612bf2fbac107722ea365932cdd35f5b\"},\"headline\":\"Training a Model on Multiple GPUs with Data Parallelism\",\"datePublished\":\"2026-01-12T14:27:19+00:00\",\"mainEntityOfPage\":{\"@id\":\"http:\/\/ktromedia.com\/?p=177895\"},\"wordCount\":1989,\"commentCount\":0,\"publisher\":{\"@id\":\"https:\/\/ktromedia.com\/#organization\"},\"image\":{\"@id\":\"http:\/\/ktromedia.com\/?p=177895#primaryimage\"},\"thumbnailUrl\":\"http:\/\/ktromedia.com\/wp-content\/uploads\/2026\/01\/Training-a-Model-on-Multiple-GPUs-with-Data-Parallelism.jpg\",\"articleSection\":[\"\u4eba\u5de5\u667a\u80fd\"],\"inLanguage\":\"en-US\",\"potentialAction\":[{\"@type\":\"CommentAction\",\"name\":\"Comment\",\"target\":[\"http:\/\/ktromedia.com\/?p=177895#respond\"]}]},{\"@type\":\"WebPage\",\"@id\":\"http:\/\/ktromedia.com\/?p=177895\",\"url\":\"http:\/\/ktromedia.com\/?p=177895\",\"name\":\"Training a Model on Multiple GPUs with Data Parallelism - Ktromedia\",\"isPartOf\":{\"@id\":\"https:\/\/ktromedia.com\/#website\"},\"primaryImageOfPage\":{\"@id\":\"http:\/\/ktromedia.com\/?p=177895#primaryimage\"},\"image\":{\"@id\":\"http:\/\/ktromedia.com\/?p=177895#primaryimage\"},\"thumbnailUrl\":\"http:\/\/ktromedia.com\/wp-content\/uploads\/2026\/01\/Training-a-Model-on-Multiple-GPUs-with-Data-Parallelism.jpg\",\"datePublished\":\"2026-01-12T14:27:19+00:00\",\"breadcrumb\":{\"@id\":\"http:\/\/ktromedia.com\/?p=177895#breadcrumb\"},\"inLanguage\":\"en-US\",\"potentialAction\":[{\"@type\":\"ReadAction\",\"target\":[\"http:\/\/ktromedia.com\/?p=177895\"]}]},{\"@type\":\"ImageObject\",\"inLanguage\":\"en-US\",\"@id\":\"http:\/\/ktromedia.com\/?p=177895#primaryimage\",\"url\":\"http:\/\/ktromedia.com\/wp-content\/uploads\/2026\/01\/Training-a-Model-on-Multiple-GPUs-with-Data-Parallelism.jpg\",\"contentUrl\":\"http:\/\/ktromedia.com\/wp-content\/uploads\/2026\/01\/Training-a-Model-on-Multiple-GPUs-with-Data-Parallelism.jpg\",\"width\":2560,\"height\":1656},{\"@type\":\"BreadcrumbList\",\"@id\":\"http:\/\/ktromedia.com\/?p=177895#breadcrumb\",\"itemListElement\":[{\"@type\":\"ListItem\",\"position\":1,\"name\":\"Home\",\"item\":\"https:\/\/ktromedia.com\/\"},{\"@type\":\"ListItem\",\"position\":2,\"name\":\"Training a Model on Multiple GPUs with Data 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":"Training a Model on Multiple GPUs with Data 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=177895","og_locale":"en_US","og_type":"article","og_title":"Training a Model on Multiple GPUs with Data Parallelism - Ktromedia","og_description":"import dataclasses import os \u00a0 import datasets import tqdm 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 from torch import Tensor from torch.nn.parallel import DistributedDataParallel as DDP from torch.utils.data.distributed import DistributedSampler \u00a0 # Build the model @dataclasses.dataclass class LlamaConfig: \u00a0\u00a0\u00a0\u00a0“”“Define Llama model hyperparameters.”“”","og_url":"http:\/\/ktromedia.com\/?p=177895","og_site_name":"Ktromedia","article_publisher":"https:\/\/www.facebook.com\/KTROMedia\/","article_published_time":"2026-01-12T14:27:19+00:00","og_image":[{"width":2560,"height":1656,"url":"http:\/\/ktromedia.com\/wp-content\/uploads\/2026\/01\/Training-a-Model-on-Multiple-GPUs-with-Data-Parallelism.jpg","type":"image\/jpeg"}],"author":"KTRO TEAM","twitter_card":"summary_large_image","twitter_misc":{"Written by":"KTRO TEAM","Est. reading time":"10 minutes"},"schema":{"@context":"https:\/\/schema.org","@graph":[{"@type":"Article","@id":"http:\/\/ktromedia.com\/?p=177895#article","isPartOf":{"@id":"http:\/\/ktromedia.com\/?p=177895"},"author":{"name":"KTRO TEAM","@id":"https:\/\/ktromedia.com\/#\/schema\/person\/612bf2fbac107722ea365932cdd35f5b"},"headline":"Training a Model on Multiple GPUs with Data Parallelism","datePublished":"2026-01-12T14:27:19+00:00","mainEntityOfPage":{"@id":"http:\/\/ktromedia.com\/?p=177895"},"wordCount":1989,"commentCount":0,"publisher":{"@id":"https:\/\/ktromedia.com\/#organization"},"image":{"@id":"http:\/\/ktromedia.com\/?p=177895#primaryimage"},"thumbnailUrl":"http:\/\/ktromedia.com\/wp-content\/uploads\/2026\/01\/Training-a-Model-on-Multiple-GPUs-with-Data-Parallelism.jpg","articleSection":["\u4eba\u5de5\u667a\u80fd"],"inLanguage":"en-US","potentialAction":[{"@type":"CommentAction","name":"Comment","target":["http:\/\/ktromedia.com\/?p=177895#respond"]}]},{"@type":"WebPage","@id":"http:\/\/ktromedia.com\/?p=177895","url":"http:\/\/ktromedia.com\/?p=177895","name":"Training a Model on Multiple GPUs with Data Parallelism - Ktromedia","isPartOf":{"@id":"https:\/\/ktromedia.com\/#website"},"primaryImageOfPage":{"@id":"http:\/\/ktromedia.com\/?p=177895#primaryimage"},"image":{"@id":"http:\/\/ktromedia.com\/?p=177895#primaryimage"},"thumbnailUrl":"http:\/\/ktromedia.com\/wp-content\/uploads\/2026\/01\/Training-a-Model-on-Multiple-GPUs-with-Data-Parallelism.jpg","datePublished":"2026-01-12T14:27:19+00:00","breadcrumb":{"@id":"http:\/\/ktromedia.com\/?p=177895#breadcrumb"},"inLanguage":"en-US","potentialAction":[{"@type":"ReadAction","target":["http:\/\/ktromedia.com\/?p=177895"]}]},{"@type":"ImageObject","inLanguage":"en-US","@id":"http:\/\/ktromedia.com\/?p=177895#primaryimage","url":"http:\/\/ktromedia.com\/wp-content\/uploads\/2026\/01\/Training-a-Model-on-Multiple-GPUs-with-Data-Parallelism.jpg","contentUrl":"http:\/\/ktromedia.com\/wp-content\/uploads\/2026\/01\/Training-a-Model-on-Multiple-GPUs-with-Data-Parallelism.jpg","width":2560,"height":1656},{"@type":"BreadcrumbList","@id":"http:\/\/ktromedia.com\/?p=177895#breadcrumb","itemListElement":[{"@type":"ListItem","position":1,"name":"Home","item":"https:\/\/ktromedia.com\/"},{"@type":"ListItem","position":2,"name":"Training a Model on Multiple GPUs with Data 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\/177895","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=177895"}],"version-history":[{"count":1,"href":"http:\/\/ktromedia.com\/index.php?rest_route=\/wp\/v2\/posts\/177895\/revisions"}],"predecessor-version":[{"id":177897,"href":"http:\/\/ktromedia.com\/index.php?rest_route=\/wp\/v2\/posts\/177895\/revisions\/177897"}],"wp:featuredmedia":[{"embeddable":true,"href":"http:\/\/ktromedia.com\/index.php?rest_route=\/wp\/v2\/media\/177896"}],"wp:attachment":[{"href":"http:\/\/ktromedia.com\/index.php?rest_route=%2Fwp%2Fv2%2Fmedia&parent=177895"}],"wp:term":[{"taxonomy":"category","embeddable":true,"href":"http:\/\/ktromedia.com\/index.php?rest_route=%2Fwp%2Fv2%2Fcategories&post=177895"},{"taxonomy":"post_tag","embeddable":true,"href":"http:\/\/ktromedia.com\/index.php?rest_route=%2Fwp%2Fv2%2Ftags&post=177895"}],"curies":[{"name":"wp","href":"https:\/\/api.w.org\/{rel}","templated":true}]}}