Pelatihan di Atas 1 Juta Token

Sebuah agen mampu mengerjakan tugas selama berjam-jam: membaca file, menjalankan perintah, dan membawa setiap giliran sesi ke depan. Sesi yang terakumulasi tersebut mencapai ratusan ribu token, itulah mengapa model frontier kini mempromosikan jendela konteks jutaan token.

Membuat model yang baik dalam menangani konteks sebanyak itu berarti melatihnya pada sekuens sepanjang itu, dan di situlah kesulitannya: sekuens satu juta token tidak muat di satu GPU, dan juga tidak muat di delapan GPU, tanpa bantuan.

Contoh dalam panduan ini melatih pada tepat satu sekuens seperti itu per langkah, pada satu node 8-GPU.

Jalankan

Anda membutuhkan satu node dengan 8 H100 (atau lebih baik), dan transformers>=5.17.0: contoh ini menggunakan offload dari gradient checkpointing, yang ditambahkan pada rilis tersebut.

torchrun --nproc_per_node 8 examples/sft_qwen3_8b_1m_context/sft_qwen3_8b_1m_context.py

Skrip ini melakukan fine-tuning Qwen3-8B pada buku dari PG-19, yang digabungkan ujung ke ujung hingga masing-masing mencapai sekitar satu juta token. Setelah sekitar sepuluh menit untuk memuat dan melakukan tokenisasi, langkah pertama tiba:

{'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]

Tiga angka yang layak dibaca:

• num_tokens adalah 1.049e6. Itu adalah satu urutan, bukan batch dari urutan-urutan pendek. Setiap token memperhatikan setiap token sebelumnya, di seluruh satu juta token.
• Loss dimulai dari 4,3, yang merupakan nilai loss awal yang normal. Jalur yang salah dikonfigurasi secara halus untuk konteks panjang dimulai sekitar 10, dan kita akan kembali membahas alasannya.
• 380 detik per langkah. Satu langkah adalah satu urutan, jadi itu sekitar enam menit lebih per urutan.

Secara naif, urutan ini membutuhkan 288 GB per GPU. Sisa panduan ini adalah cerita tentang bagaimana angka itu turun menjadi 56.

Apa yang baru saja terjadi

Semua di bawah ini dibangun dari setup yang sudah kita kenal: satu GPU, beberapa ribu token per sekuens. Masalah mulai muncul seiring pertumbuhan sekuens, dan masalah tersebut muncul dalam urutan yang tetap.

Loss adalah hal pertama yang rusak

Ambil Qwen3-4B di satu kartu 80 GB dan tingkatkan panjang sekuens. Hal pertama yang kehabisan memori bukanlah modelnya. Melainkan loss-nya. Periksa profil memori.

Puncak terjadi di antara langkah maju dan langkah mundur, yaitu tempat di mana loss dihitung. Untuk memahami mengapa puncak berada di sana, kita perlu melihat apa yang dilakukan oleh tahap akhir dari decoding dan bagaimana loss dibangun dari tahap-tahap tersebut.

Lapisan terakhir dari decoder menghasilkan hidden state dengan bentuk$L \times H$, di mana$L$ adalah panjang sekuens dan$H$ adalah ukuran hidden. Hidden state ini kemudian dikalikan dengan bobot kepala pemodelan bahasa untuk menghasilkan matriks logits dengan bentuk$L \times V$, di mana$V$ adalah ukuran kosakata. Matriks ini sangat besar karena sekuensnya panjang dan kosakatanya besar (lebih dari seratus ribu entri). Loss kemudian dihitung dengan fungsi cross entropy dari matriks tersebut dan label target, dan materialisasi matriks inilah yang menjadi puncak penggunaan memori.

Bisakah kita menghitung loss tanpa mematerialisasi seluruh matriks logits? Ya, kita dapat memprosesnya dalam chunk.

Matriks logits yang masif: pecah menjadi beberapa bagian

Alih-alih menghitung logits untuk seluruh rangkaian sekaligus, kita mengambil beberapa ratus baris pada satu waktu, menghitung loss untuk baris-baris tersebut, lalu menjumlahkan bagian-bagiannya. Matriks logits secara keseluruhan tidak pernah disimpan di memori sekaligus. Baris adalah dimensi yang terbagi dengan bersih: loss adalah jumlahan atas baris, dan softmax di dalamnya berjalan sepanjang kosakata dari satu baris, sehingga sepotong baris membawa semua yang dibutuhkan loss-nya sendiri.

Satu perkalian besar menjadi beberapa perkalian yang lebih kecil, dan penghematan memori cukup besar sehingga kerugian tidak lagi menjadi puncak. TRL menggunakan 256 baris per chunk, yang menjaga setiap perkalian cukup besar untuk menggunakan GPU dengan baik.

Profil baru terlihat seperti ini.

Puncaknya telah hilang.

Bagaimana cara mengaktifkan ini di TRL? Tidak ada yang perlu dilakukan: chunked loss adalah default. Jika Anda ingin menggunakan plain loss, set loss_type="nll" di SFTConfig.

Sejauh mana hal ini memungkinkan kita mendorong sekuens? Dalam hal memori, perubahan tunggal ini saja membawa kita dari 32k token menjadi 160k.

Model belum pernah melihat posisi 100.000: skala ulang posisi dengan YaRN

Sekarang kita bisa memuat 160k token ke dalam memori. Tapi apakah model benar-benar bisa membacanya?

Qwen3-4B dilatih pada urutan 40.960 token. Konfigurasinya menyatakan demikian:

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

Setiap posisi di luar itu adalah posisi yang belum pernah ditemui oleh model. Tidak ada yang crash, tetapi proses hanya melatih pada token yang tidak dapat ditempatkan oleh model. Berikut adalah loss pada setiap token dari sekuens 160k-token:

Model tidak runtuh pada 40.960. Ia terus berjalan untuk sementara waktu, dan baru mulai kehilangan alur di sekitar 70 ribu. Dari titik itu, nilainya naik secara stabil, berakhir di atas entropi unigram dari teks yang sama (6,45): lebih buruk daripada model yang mengabaikan konteks sepenuhnya dan hanya memprediksi berdasarkan frekuensi token.

Setiap token memiliki posisinya, dan model telah mempelajari arti posisi antara 0 dan 40.960. Jika Anda bertanya tentang posisi 100.000, model tidak memiliki dasar untuk menjawab.

Jalan keluarnya adalah jangan pernah memberikan posisi yang tidak dikenalnya. Kita melakukan penskalaan ulang: posisi 100.000 menjadi posisi 25.000, dan setiap posisi dalam urutan kembali berada di dalam rentang yang diketahui model. Token tidak berubah, hanya posisi yang diberikan untuk masing-masing token. Inilah yang dilakukan oleh YaRN, dan faktornya adalah seberapa jauh kita ingin pergi: 40.960 terlatih, 163.840 diinginkan, sehingga 4.

Bagaimana cara mengaktifkan ini di TRL? Sampaikan konfigurasi RoPE saat model dimuat:

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,
        },
    },
)

Pengukuran yang sama, dengan run yang diskalakan ulang di bagian atas:

Jalankan yang diskalakan ulang tetap datar hingga 160k. Pada 20k token terakhir, rata-rata loss-nya adalah 2,8, dibandingkan dengan 7,3 tanpa penskalaan ulang.

Aktivasi adalah yang tersisa: offload-kan

Dengan loss yang di-chunk dan posisi yang di-rescale, satu GPU membawa kita hingga 160k token. Lalu kehabisan lagi. Apa yang mengisi kartu kali ini?

Model dan optimizer-nya menyumbang 22 GB tetap, apa pun panjang sekuensnya. Di atas itu, setiap pita merah adalah aktivasi yang disimpan dari satu lapisan, masing-masing 0,47 GB. Pita-pita ini menumpuk saat forward menuliskannya, diam di sana, dan dilepas satu per satu saat backward mengonsumsinya. Bagian abu-abu di atas adalah semua yang dihitung ulang dan dibuang dalam hitungan milidetik.

TRL sudah mengurangi mereka untuk Anda, dan ini diaktifkan secara default. Gradient checkpointing hanya menyimpan apa yang masuk ke setiap layer dan membuang semua yang dihitung layer di dalam dirinya: skor attention, intermediate MLP yang lebar, dan norm. Ketika backward pass membutuhkannya, layer tersebut dijalankan untuk kedua kalinya untuk memperolehnya kembali. Yang tersisa adalah satu tensor tersimpan per layer, masing-masing berukuran sequence x hidden, dan pada panjang sequence yang panjang, itu masih merupakan hal terbesar di kartu.

Pengamatan yang berguna adalah bahwa tensor-tensor tersebut sama sekali tidak diperlukan selama forward pass. Mereka ditulis, dibiarkan, dan dibaca jauh kemudian. Jadi, mereka tidak perlu berada di GPU di antara waktu tersebut: mereka dapat dipindahkan ke memori CPU dan dikembalikan ketika backward pass mencapai mereka.

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

Langkah yang sama, panjang yang sama: alih-alih 38 band yang menumpuk, hanya 4 yang pernah hadir sekaligus, dan puncak turun dari 59,9 GB menjadi 48,8 GB. Hal ini memungkinkan panjang sekuens:

Di bawah 48k, keduanya sama, karena pada panjang tersebut hampir tidak ada yang perlu dipindahkan. Melewati titik itu, eksekusi yang di-offload tetap lebih rendah, dan di mana urutan sebelumnya berhenti di 160k, kini mencapai 256k pada kartu yang sama.

Yang satu ini tidak gratis. Setiap aktivasi yang disimpan berpindah ke host dan kembali, lalu lalu lintas tersebut harus muat di antara komputasi yang berada di sebelahnya.

Satu GPU tidak cukup: bagi urutan ke banyak

Chunked loss, offloaded activations: satu kartu kini melatih pada 256k token. Bagaimana mencapai satu juta? Untuk melihat apa yang menghalangi, lihat ke dalam satu lapisan selama backward pass (puncak abu-abu pada profil di atas).

Gradient checkpointing berarti forward hampir tidak menyimpan apa pun: setiap layer menyimpan inputnya dan membuang sisanya. Layer tersebut dihitung ulang nanti, ketika backward mencapainya, dan itulah saat di mana semuanya ada sekaligus. Di dalam MLP, satu baris dari model,

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

empat tensor dibangun pada lebar intermediat penuh, dan implementasi standar mempertahankan keempatnya hingga backward melewati mereka. Pada satu juta token, pada model yang melebarkan setiap token dari 2560 menjadi 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, di dalam satu lapisan, pada kartu yang memiliki 80 dan sudah memberikan 30 untuk model dan optimizer-nya. Pada titik ini, jawabannya bukan untuk menyimpan urutan dengan cara yang lebih cerdas. Jawabannya adalah memberikan lebih sedikit dari itu ke setiap GPU.

Beri setiap GPU satu potongan dari sekuens. Dengan empat GPU, setiap GPU menyimpan 262.144 token alih-alih 1.048.576. Keempat tensor menyusut bersamanya: 19 GB alih-alih 76, yang bersama dengan 30 GB model muat di kartu.

Keberatan yang jelas adalah perhatian. Setiap token harus memperhatikan setiap token sebelumnya, dan tidak ada GPU yang memiliki token-token sebelumnya tersebut. Jadi, GPU-gpu tersebut saling bertukar apa yang mereka kekurangan, dan hasilnya persis sama dengan apa yang akan dihitung oleh satu GPU yang sangat besar.

Ada dua cara untuk menjalankan pertukaran tersebut, yaitu paralelisme konteks (CP) dan paralelisme urutan Ulysses (SP).

Pada FSDP2, gunakan CP. Setiap GPU menyimpan slice-nya sendiri, dan slice lainnya datang secara bergiliran, sehingga setiap token pada akhirnya melihat setiap token sebelumnya.

from accelerate import ParallelismConfig

training_args = SFTConfig(..., fsdp=True, parallelism_config=ParallelismConfig(cp_size=4))

Pada DeepSpeed, gunakan SP. Alih-alih memindahkan potongan, SP menyusun ulang batch tepat sebelum attention sehingga setiap GPU memegang semua token tetapi hanya seperempat dari kepala attention, lalu menyusun ulang kembali setelahnya.

training_args = SFTConfig(
    ..., deepspeed="examples/deepspeed_configs/zero3.json", parallelism_config=ParallelismConfig(sp_size=4, sp_backend="deepspeed")
)

Dua pengaturan ini tidak saling tumpang tindih saat ini: cp_size memerlukan FSDP2 dan sp_size hanya berjalan di bawah DeepSpeed. Pilih backend-nya, dan metodenya akan mengikuti. Sisa bagian ini mengikuti jalur FSDP2, yang merupakan jalur yang digunakan oleh contoh di bagian atas panduan ini: GPU-nya bukan pekerja independen pada batch terpisah, melainkan satu grup yang berbagi satu urutan tunggal.

Dua hal berubah setelahnya diaktifkan:

1. Sequences harus di-pad hingga kelipatan dari cp_size * 2, jadi pad_to_multiple_of=8 (di dalam SFTConfig) untuk empat GPU.
2. Kewajiban causal-SDPA menyingkirkan packing, yang bergantung pada maska blok-diagonal untuk mencegah dokumen saling membaca, dan TRL akan memicu error jika Anda meminta keduanya. Hal ini juga menyingkirkan model-model yang lapisannya menggunakan sliding-window atau chunked attention, yang ditolak oleh accelerate: OpenAI GPT-OSS, Gemma 3 dan 4, Qwen3.5 dan seterusnya. Qwen3 dan Qwen3-MoE menggunakan full attention di seluruh lapisan, itulah mengapa mereka adalah model-model di sini.

Mengirimkan potongan data ke berbagai tempat memiliki biaya yang lebih rendah dari yang mungkin Anda perkirakan. Pada 131k token di satu node, satu langkah membutuhkan waktu 34,6 detik di dua GPU, 17,8 detik di empat, dan 9,5 detik di delapan. Setiap penggandaan kelompok hampir mengurangi waktu langkah menjadi setengahnya, dan perpindahan dari dua ke delapan GPU memulihkan 3,7x dari potensi 4x.

Dan itulah tuas terakhir dari empat tuas:

Satu kartu berhenti tepat di atas 256k. Empat kartu menangani satu juta penuh!, mendarat di 63,6 GB dengan ruang cadangan. Urutan yang membuka panduan ini muat.

Set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True. Pada satu juta token, eksekusi di atas membutuhkan 63,6 GB pada kartu 80 GB, dan tetap gagal tanpa pengaturan ini: tensor yang diminta berukuran besar dan berurutan, sehingga allocator tidak dapat menemukan ruang untuknya di antara blok yang sudah dipegang.

Bacaan lanjutan

• Paralelisme konteks, panduan accelerate untuk cp_size dan jaring perangkat yang dibangunnya.
• Ulysses dan ring attention, tentang bagaimana kedua pertukaran tersebut berbeda dan berapa biaya komunikasi masing-masing.
• YaRN, penyesuaian ulang posisi yang digunakan di bagian RoPE.