Huấn luyện vượt quá 1 triệu token

Một agent có khả năng thực hiện một tác vụ trong nhiều giờ: đọc tệp, chạy lệnh và duy trì toàn bộ các lượt của phiên làm việc. Phiên tích lũy đó có thể lên đến hàng trăm nghìn token, đó là lý do tại sao các mô hình tiên phong hiện nay quảng bá các cửa sổ ngữ cảnh lên tới một triệu token.

Việc giúp một mô hình xử lý tốt lượng ngữ cảnh lớn như vậy đòi hỏi phải huấn luyện nó trên các chuỗi có độ dài tương đương, và đây chính là nơi trở nên khó khăn: một chuỗi một triệu token không thể chứa vừa trên một GPU, và cũng không thể chứa vừa trên tám GPU, nếu không có sự hỗ trợ.

Ví dụ trong hướng dẫn này huấn luyện trên chính xác một chuỗi như vậy mỗi bước, trên một nút duy nhất có 8 GPU.

Chạy

Bạn cần một node có 8 H100 (hoặc tốt hơn), và transformers>=5.17.0: ví dụ sử dụng offload của gradient checkpointing, được thêm vào trong bản phát hành đó.

accelerate launch \
    --config_file examples/sft_qwen3_8b_1m_context/context_parallel_8gpu.yaml \
    examples/sft_qwen3_8b_1m_context/sft_qwen3_8b_1m_context.py

Script này fine-tune Qwen3-8B trên các cuốn sách từ PG-19, nối liền nhau cho đến khi mỗi cuốn dài khoảng một triệu token. Sau khoảng mười phút để tải và tokenize, bước đầu tiên được thực hiện:

{'loss': 4.311, 'grad_norm': 29.25, 'num_tokens': 1049000.0, 'epoch': 0.25}
  8%|▊ | 1/12 [06:20<1:09:41, 380.10s/it]

Ba con số đáng chú ý:

• num_tokens là 1.049e6. Đó là một chuỗi duy nhất, không phải một lô gồm các chuỗi ngắn. Mỗi token chú ý đến mọi token trước đó, trải dài trên toàn bộ một triệu token.
• Loss bắt đầu ở mức 4.3, đây là mức loss khởi đầu bình thường. Một lần chạy bị cấu hình sai một cách tinh vi cho ngữ cảnh dài sẽ bắt đầu ở khoảng 10, và chúng ta sẽ quay lại lý do tại sao.
• 380 giây mỗi bước. Một bước là một chuỗi, vì vậy đó là hơn sáu phút mỗi chuỗi.

Một cách ngây thơ, chuỗi này cần 288 GB cho mỗi GPU. Phần còn lại của hướng dẫn này là câu chuyện về cách con số đó giảm xuống còn 56.

Điều vừa xảy ra

Mọi thứ bên dưới được xây dựng dựa trên một cấu hình mà chúng ta đã quen thuộc: một GPU, vài nghìn token cho mỗi chuỗi. Các vấn đề bắt đầu xuất hiện khi chuỗi dài hơn, và chúng xảy ra theo một thứ tự cố định.

Mất mát là điều đầu tiên bị hỏng

Hãy thử Qwen3-4B trên một card 80 GB duy nhất và tăng độ dài chuỗi. Thứ đầu tiên hết bộ nhớ không phải là mô hình. Đó là loss. Kiểm tra hồ sơ bộ nhớ.

Đỉnh nằm giữa bước truyền thuận và bước truyền ngược, nơi mà loss được tính toán. Để hiểu tại sao đỉnh lại nằm ở đó, chúng ta cần xem xét các giai đoạn cuối của quá trình giải mã thực hiện gì và loss được xây dựng từ chúng như thế nào.

Lớp cuối cùng của decoder xuất ra một hidden state có hình dạng$L \times H$, trong đó$L$ là độ dài chuỗi và$H$ là kích thước ẩn. Hidden state này sau đó được nhân với trọng số của đầu ra mô hình hóa ngôn ngữ để tạo ra một ma trận logits có hình dạng$L \times V$, trong đó$V$ là kích thước từ vựng. Ma trận này rất lớn vì chuỗi dài và từ vựng lớn (hơn một trăm nghìn mục). Sau đó, loss được tính bằng hàm cross entropy từ ma trận đó và các nhãn đích, và việc vật chất hóa nó chính là nguyên nhân gây ra đỉnh.

Chúng ta có thể tính toán loss mà không cần tạo ra toàn bộ ma trận logits không? Có, chúng ta có thể xử lý theo từng khối.

Ma trận logits khổng lồ: chia nhỏ

Thay vì tính toán logits cho toàn bộ chuỗi một lúc, chúng ta lấy vài trăm dòng một lần, tính loss cho các dòng đó, rồi cộng các phần lại với nhau. Ma trận logits hoàn chỉnh không bao giờ được giữ trong bộ nhớ cùng một lúc. Dòng là chiều có thể tách rời một cách sạch sẽ: loss là tổng trên các dòng, và softmax bên trong nó chạy dọc theo từ vựng của một dòng, do đó một khối các dòng mang theo mọi thứ mà loss của chính nó cần.

Một phép nhân lớn được chia thành nhiều phép nhân nhỏ hơn, và mức tiết kiệm bộ nhớ đủ lớn để phần mất mát không còn là đỉnh nữa. TRL sử dụng 256 hàng cho mỗi khối, giúp mỗi phép nhân đủ lớn để khai thác GPU hiệu quả.

Hồ sơ mới trông như thế này.

Đỉnh đã biến mất.

Làm thế nào để bật điều này trong TRL? Không cần làm gì: chunked loss là mặc định. Nếu bạn muốn chọn plain loss, hãy đặt loss_type="nll" trong SFTConfig.

Điều này cho phép chúng ta đẩy dài chuỗi đến mức nào? Về mặt bộ nhớ, chỉ riêng thay đổi này đã giúp chúng ta tăng từ 32k token lên 160k.

Mô hình chưa bao giờ thấy vị trí 100.000: hãy điều chỉnh lại các vị trí bằng YaRN

Bây giờ chúng ta có thể nạp 160k token vào bộ nhớ. Nhưng liệu mô hình có thực sự đọc được chúng không?

Qwen3-4B được huấn luyện trên các chuỗi gồm 40.960 token. Cấu hình xác nhận điều này:

>>> from transformers import AutoConfig
>>> AutoConfig.from_pretrained("Qwen/Qwen3-4B").max_position_embeddings
40960

Mọi vị trí vượt quá ngưỡng đó đều là những vị trí mà mô hình chưa từng gặp. Không có gì bị lỗi, nhưng quá trình chạy chỉ huấn luyện trên các token mà mô hình không thể xác định vị trí. Dưới đây là loss trên từng token của một chuỗi 160k token:

Mô hình không sụp đổ ở mức 40.960. Nó tiếp tục hoạt động một thời gian và chỉ bắt đầu mất mạch sau khoảng 70k. Từ đó, nó tăng đều đặn, kết thúc ở mức cao hơn entropy unigram của cùng văn bản đó (6.45): tệ hơn một mô hình bỏ qua hoàn toàn ngữ cảnh và chỉ dự đoán dựa trên tần suất token.

Mỗi token đi kèm với vị trí của nó, và mô hình đã học được ý nghĩa của các vị trí nằm trong khoảng từ 0 đến 40.960. Nếu hỏi về vị trí 100.000, mô hình sẽ không có cơ sở nào để suy luận.

Cách thoát ra là không bao giờ đưa cho nó một vị trí mà nó không biết. Chúng ta sẽ chuẩn hóa lại: vị trí 100.000 trở thành vị trí 25.000, và mọi vị trí trong chuỗi đều quay lại nằm trong khoảng mà mô hình biết. Các token không thay đổi, chỉ có vị trí được gán cho từng token là thay đổi. Đây là điều mà YaRN làm, và hệ số đơn giản là khoảng cách chúng ta muốn đi xa hơn: 40.960 đã được huấn luyện, 163.840 là mục tiêu, nên hệ số là 4.

Làm thế nào để bật điều này trong TRL? Truyền cấu hình RoPE khi tải mô hình:

training_args = SFTConfig(
    ...,
    model_init_kwargs={
        "rope_parameters": {
            "rope_type": "yarn",
            "rope_theta": 1_000_000,  # has to match what the model ships with
            "factor": 4.0,
            "original_max_position_embeddings": 40960,
        },
    },
)

Cùng một phép đo, với lần chạy đã được điều chỉnh tỷ lệ ở trên:

Chạy được điều chỉnh tỷ lệ giữ nguyên mức phẳng cho đến 160k. Trong 20k token cuối cùng, nó đạt trung bình loss là 2.8, so với 7.3 khi không điều chỉnh tỷ lệ.

Các activation là thứ còn lại: hãy offload chúng

Với việc loss được chia nhỏ và các vị trí được điều chỉnh lại, một GPU có thể xử lý tới 160k token. Sau đó, bộ nhớ lại đầy. Lần này, thứ gì đang chiếm hết card?

Mô hình và bộ tối ưu hóa của nó chiếm một dung lượng cố định là 22 GB, bất kể độ dài chuỗi. Ngoài phần đó, mỗi dải màu đỏ là một phần kích hoạt được lưu của một lớp, mỗi phần 0,47 GB. Chúng được xếp chồng lên nhau khi quá trình truyền thẳng ghi chúng, nằm đó, và được gỡ bỏ lần lượt khi quá trình truyền ngược sử dụng chúng. Phần màu xám ở trên cùng là tất cả những gì được tính lại và bị loại bỏ trong vòng vài mili giây.

TRL đã tự động giảm thiểu chúng cho bạn, và tính năng này được bật theo mặc định. Gradient checkpointing chỉ giữ lại những gì đi vào mỗi lớp và loại bỏ mọi thứ mà lớp tính toán bên trong: các điểm chú ý (attention scores), phần trung gian rộng của MLP, các phép chuẩn hóa (norms). Khi quá trình truyền ngược (backward pass) cần đến chúng, nó sẽ chạy lại lớp đó một lần nữa để lấy lại. Những gì còn lại là một tensor được lưu cho mỗi lớp, mỗi tensor có kích thước sequence x hidden, và ở độ dài chuỗi dài, đây vẫn là thành phần lớn nhất trên card.

Quan sát hữu ích là các tensor đó hoàn toàn không cần thiết trong quá trình truyền thẳng. Chúng được ghi, để yên, và chỉ được đọc lại sau đó rất lâu. Do đó, trong khoảng thời gian đó, chúng không cần phải nằm trên GPU: chúng có thể được chuyển sang bộ nhớ CPU và đưa trở lại khi quá trình truyền ngược đến chúng.

training_args = SFTConfig(..., gradient_checkpointing_kwargs={"offload": True})

Cùng bước, cùng độ dài: thay vì 38 băng xếp chồng lên nhau, chỉ có 4 băng hiện diện cùng lúc, và đỉnh giảm từ 59,9 GB xuống 48,8 GB. Điều này đổi lấy độ dài chuỗi:

Dưới 48k, hai phương pháp cho kết quả như nhau vì ở độ dài đó hầu như không có gì để di chuyển. Vượt qua ngưỡng đó, phiên bản offload duy trì mức thấp hơn, và nơi chuỗi trước đây dừng lại ở 160k thì nay đạt tới 256k trên cùng một card.

Phương pháp này không miễn phí. Mỗi activation được lưu trữ đều phải di chuyển sang host và quay lại, và lưu lượng đó phải vừa vặn với phần tính toán nằm cạnh nó.

Một GPU là không đủ: chia nhỏ chuỗi trên nhiều

Chunked loss, offloaded activations: một card giờ có thể huấn luyện trên 256k token. Làm thế nào để đạt tới một triệu? Để xem điều gì đang cản trở, hãy nhìn vào bên trong một lớp duy nhất trong quá trình backward (các đỉnh màu xám trong hồ sơ ở trên).

Gradient checkpointing có nghĩa là bước forward gần như không lưu trữ gì: mỗi lớp chỉ lưu đầu vào của nó và loại bỏ phần còn lại. Lớp đó sẽ được tính lại sau, khi bước backward đến nó, và đó là khoảnh khắc mọi thứ cùng tồn tại. Trong MLP, một dòng của mô hình,

down_proj(act_fn(gate_proj(x)) * up_proj(x))

bốn tensor được xây dựng ở độ rộng trung gian đầy đủ, và cài đặt chuẩn giữ lại cả bốn cho đến khi quá trình backward đi qua chúng. Với một triệu token, trên một mô hình mà MLP mở rộng mỗi token từ 2560 lên 9728:

gate_proj(x)                             19 GB = 1,048,576 x 9728 x 2 bytes
act_fn(gate_proj(x))                     19 GB
up_proj(x)                               19 GB
act_fn(gate_proj(x)) * up_proj(x)        19 GB
                                  total  76 GB

76 GB, nằm trong một lớp, trên một card có 80 và đã cấp 30 cho model cùng optimizer của nó. Đến thời điểm này, câu trả lời không phải là lưu trữ sequence một cách thông minh hơn. Mà là cấp ít hơn cho mỗi GPU.

Gán cho mỗi GPU một lát của chuỗi. Với bốn GPU, mỗi GPU chỉ giữ 262.144 token thay vì 1.048.576. Bốn tensor co lại theo: 19 GB thay vì 76, cùng với 30 GB của mô hình, vừa vặn trên card.

Lý do phản đối rõ ràng nhất là cơ chế chú ý. Mỗi token phải chú ý đến mọi token trước đó, và không có GPU nào nắm giữ các token trước đó. Do đó, các GPU trao đổi cho nhau những phần còn thiếu, và kết quả thu được chính xác là những gì một GPU khổng lồ duy nhất sẽ tính toán.

Có hai cách để thực hiện quá trình trao đổi đó: song song ngữ cảnh (CP) và song song chuỗi Ulysses (SP).

Trên FSDP2, hãy sử dụng CP. Mỗi GPU giữ một lát dữ liệu của riêng mình, và các lát dữ liệu khác lần lượt đến với nó, nhờ đó mỗi token cuối cùng đều nhìn thấy mọi token trước đó.

parallelism_config:
  parallelism_config_cp_size: 4

Trên DeepSpeed, hãy dùng SP. Thay vì di chuyển các lát dữ liệu, nó sắp xếp lại batch ngay trước khi tính attention để mỗi GPU giữ mọi token nhưng chỉ một phần tư các đầu attention, sau đó xáo trộn lại như cũ.

parallelism_config:
  parallelism_config_sp_size: 4

Hai tùy chọn này hiện không chồng chéo lên nhau: cp_size yêu cầu FSDP2 và sp_size chỉ chạy dưới DeepSpeed. Bạn chọn backend và phương pháp sẽ theo sau. Phần còn lại của mục này đi theo đường dẫn FSDP2, đây là cách mà ví dụ ở đầu hướng dẫn này sử dụng: các GPU của nó không phải là những worker độc lập trên các batch riêng biệt, mà là một nhóm chia sẻ một chuỗi duy nhất.

Hai điều thay đổi khi nó được bật:

1. Các chuỗi phải được đệm (padding) đến bội số của cp_size * 2, do đó pad_to_multiple_of=8 (trong SFTConfig) cho bốn GPU.
2. Yêu cầu causal-SDPA loại trừ việc đóng gói (packing), vốn dựa vào mặt nạ khối đường chéo (block-diagonal mask) để ngăn các tài liệu đọc lẫn nhau, và TRL sẽ báo lỗi nếu bạn yêu cầu cả hai. Nó cũng loại trừ các mô hình có các lớp sử dụng chú ý cửa sổ trượt (sliding-window) hoặc chú ý phân mảnh (chunked attention), những thứ mà accelerate từ chối: OpenAI GPT-OSS, Gemma 3 và 4, Qwen3.5 và các phiên bản sau. Qwen3 và Qwen3-MoE sử dụng chú ý đầy đủ (full attention) xuyên suốt, đây là lý do chúng là các mô hình ở đây.

Việc truyền các lát dữ liệu qua lại có chi phí thấp hơn bạn có thể tưởng tượng. Với 131k token trên một node, một bước (step) mất 34,6 giây trên hai GPU, 17,8 giây trên bốn và 9,5 giây trên tám. Mỗi lần nhân đôi nhóm gần như giảm một nửa thời gian mỗi bước, và việc tăng từ hai lên tám GPU đã khôi phục 3,7 lần so với mức tối đa có thể là 4 lần.

Và đó là đòn bẩy thứ tư và cũng là cuối cùng:

Một thẻ dừng lại ngay sau 256k. Bốn thẻ xử lý trọn vẹn một triệu!, đạt 63,6 GB với dư địa thoải mái. Chuỗi mà hướng dẫn này mở đầu sẽ vừa vặn.

Đặt PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. Với một triệu token, lượt chạy ở trên cần 63,6 GB trên một card 80 GB, và vẫn thất bại nếu không có nó: các tensor mà nó yêu cầu có kích thước lớn và liên tục, và bộ cấp phát bộ nhớ không thể tìm được chỗ trống cho chúng giữa các khối mà nó đã giữ.

Đọc thêm

• Song song ngữ cảnh, hướng dẫn của accelerate về cp_size và lưới thiết bị mà nó xây dựng.
• Ulysses và ring attention, về cách hai kiểu trao đổi khác nhau và chi phí truyền thông của từng kiểu.
• YaRN, phương pháp điều chỉnh lại vị trí được sử dụng trong phần RoPE.