Tính lại biểu diễn trung gian để tiết kiệm memory training

1Định nghĩa

Activation checkpointing (hay gradient checkpointing) là kỹ thuật đổi tính toán lấy bộ nhớ: thay vì giữ mọi giá trị trung gian (activation) của lượt forward cho tới khi cần, hệ thống chỉ giữ vector đầu vào ở ranh giới mỗi lớp Transformer, xóa các tensor bên trong lớp, rồi chạy lại forward của lớp đó khi backward lùi tới. Mặc định, lượng bộ nhớ HBM dành cho activation tăng tuyến tính theo số lớp, cỡ batch và độ dài chuỗi, vì tất cả phải chờ trên GPU tới lượt backward lấy ra tính đạo hàm.

Khi backward lùi tới lớp nào, runtime lấy đầu vào ranh giới đã cất để chạy lại nhanh forward của riêng lớp đó rồi dùng kết quả tính gradient. Cái giá là thêm một lượt forward (khoảng một phần ba thời gian cả bước theo ước tính ở bài này); đổi lại bộ nhớ activation giảm mạnh, cho phép tăng batch hoặc độ dài ngữ cảnh. Kỹ thuật này tương thích và thường dùng cùng với ZeRO.

Ví dụ đời thường: Activation checkpointing giống cách làm bài toán dài khi mặt bàn quá chật không đủ chỗ trải 20 tờ giấy nháp: bạn chỉ giữ tờ ghi kết quả cuối mỗi chương lớn và cất các tờ nháp lẻ đi; lúc kiểm tra ngược lại tới chương nào, bạn nháp lại đúng chương đó từ tờ mốc đã giữ. Cái giá là mất thêm thời gian nháp lại, và nếu chương có bước dựa vào may rủi (như tung xúc xắc) thì phải ghi lại kết quả tung trước đó.

2Luồng

Đây là ô 4B trong cột 4 · Scale & phương pháp. Giống ô 4A, nó không đổi phép toán học mà đổi nơi giữ dữ liệu: mũi tên kéo dài từ forward (cột 2) sang backward (cột 3), và kỹ thuật này đổi một phần mũi tên lưu trữ thành mũi tên tính lại.

Sơ đồ kiến trúc huấn luyện · 22 ô (1A → 5C)

Kiến trúc huấn luyện LLM · Ánh xạ với Bài Lab tổng hợp

Thứ tự đọc: 1A → 1B ở cột 1, 2A → 2D ở cột 2, 3A → 3E ở cột 3; mũi tên next batch từ 3E quay về cột 1 và mũi tên load weights đi xuống băng 5 (5A → 5C). Cột 4 không phải bước kế tiếp: 4A, 4B tiết kiệm bộ nhớ, còn 4C–4G là các phương pháp huấn luyện thay thế nhau. Di chuột hoặc bấm vào từng ô để đọc pop-over.

Tất cả (22)1. Dữ liệu & nhãn2. Forward3. Học từ sai số4. Quy mô & phương pháp5. Suy luận

Ẩn khung overlay

1A1B1C2A2B2C2D3A3B3C3D3E4A4B4C4D4E4F4G5A5B5C

4B4 · Scale & Training Methods

drop activations in forward → recompute in backward · memory ↓ compute ↑

4B · Activation Recompute

Bỏ bớt kết quả trung gian của forward rồi tính lại khi backward cần, đổi thêm phép tính lấy bộ nhớ.

Nhận gì

Activation (các kết quả trung gian của forward) mà backward ở ô 3B cần dùng lại.

Làm gì, vì sao

Thay vì giữ mọi activation đến lúc backward, bỏ phần lớn chúng trong forward rồi tính lại từng đoạn khi backward cần. Đổi thêm một lượt tính lấy mức giảm lớn bộ nhớ, vì activation có thể lớn hơn cả trọng số.

Ra gì

Cùng gradient như khi lưu đầy đủ (nếu phép tính xác định), nhưng đỉnh bộ nhớ thấp hơn. Kết quả đi về ô 3B.

Ví dụ, cách đọc ô

Đọc nhãn: drop activations in forward → recompute in backward; memory ↓ · compute ↑. Ví dụ đơn vị tự đặt: 32 lớp, mỗi lớp 10 đơn vị activation, giữ hết là 320, checkpoint từng lớp còn 42, đổi lại thêm khoảng 33% phép tính.

Lưu ý

Không nhầm với checkpoint lưu tệp ở ô 3E: recompute diễn ra ở mỗi bước và không tạo tệp. Nó giảm bộ nhớ đỉnh chứ không chắc làm mỗi bước nhanh hơn.

Đọc bài 4B · Activation Recompute↗Giai đoạn 8 trong Bài Lab: Bộ nhớ: ZeRO & recompute↗

flowchart TD
  step1["Forward đoạn được chọn"]
  step2["Bỏ giá trị trung gian bên trong"]
  step3["Backward cần giá trị"]
  step4["Tính gradient như thường"]
  step1 --> step2
  step2 --> step3
  step3 --> step4

3Giải pháp

Vị trí trên sơ đồ: ô 4B "Activation recompute" trong cột 4 · Scale & phương pháp. Nó nằm trên đường nối giữa forward (cột 2) và backward (ô 3B), và cùng với ô 4A (ZeRO) giải bài toán bộ nhớ nhưng ở phần khác: ZeRO chia model states, ô này giảm activation.

Vấn đề cần giải. Backward cần các giá trị trung gian của forward, nên bình thường mọi activation của mọi lớp phải nằm trên GPU cho đến khi backward dùng. Với mô hình sâu và chuỗi dài, activation có thể lớn hơn cả trọng số.

Khái niệm tối thiểu. Activation là đầu ra của từng phép trong forward. Ranh giới (boundary) là tensor đầu vào của một đoạn. Nếu giữ ranh giới và còn nguyên trọng số, ta có thể chạy lại forward của đoạn để lấy lại mọi activation bên trong. Tính lại không đổi kết quả toán học nếu phép tính xác định (hoặc nguồn ngẫu nhiên được khôi phục).

Ví dụ có số (minh họa, mô hình đơn vị tự đặt). Mô hình 32 lớp. Giả sử mỗi lớp tạo 10 đơn vị activation bên trong và đầu vào ranh giới chiếm 1 đơn vị. Không checkpoint: 32 × 10 = 320 đơn vị. Checkpoint từng lớp: giữ 32 ranh giới = 32 đơn vị cộng 10 đơn vị cho lớp đang được tính lại = 42 đơn vị, bằng khoảng 13% so với 320.

Chi phí tính toán, cũng theo mô hình đơn giản. Gọi thời gian forward là F và backward thường xấp xỉ 2F. Bình thường một bước là F + 2F = 3F. Với checkpoint từng lớp có thêm một lượt forward chạy lại: F + F + 2F = 4F, tức tăng khoảng 33%. Đây là phép ước tính từ giả định, thời gian thật phụ thuộc kernel, kích thước lớp và phần forward bị bỏ qua (ví dụ phần không cần tính lại).

Cách đọc. Ví dụ cho thấy hai thái cực: giữ hết thì tốn bộ nhớ, bỏ hết (chỉ giữ đầu vào mô hình) thì phải tính lại nhiều. Thực tế chọn độ hạt (granularity): checkpoint từng lớp, từng nhóm lớp hoặc chỉ một số khối đắt bộ nhớ như attention. Kết quả không thể kết luận: giảm peak memory chưa nói gì về tốc độ huấn luyện cả bước hay chất lượng mô hình.

Thực thi. Dùng API checkpoint của framework để nó giữ đúng phụ thuộc và khôi phục trạng thái ngẫu nhiên khi cần. Đừng chỉ đo mức đỉnh bộ nhớ giảm: nếu bước chậm đi quá nhiều, tổng năng suất huấn luyện có thể xấu hơn. Hãy so gradient với đường không tính lại, vì sai khác nằm ở sai số số học chứ không ở toán học. Đừng nhầm với checkpoint lưu tệp ở ô 3E (vòng lặp, eval, checkpoint): activation checkpointing diễn ra mỗi bước và không tạo tệp mô hình.

Để kiểm chứng tính lại cho ra cùng gradient với đường lưu đầy đủ, mở giai đoạn 8 · Bộ nhớ: ZeRO & recompute của Bài Lab tổng hợp. Lưu ý phép kiểm tra đó nói về sự trùng khớp số học, không đo bộ nhớ đỉnh thật. Hãy đoán trước sai khác lớn nhất giữa hai đường gradient là cỡ nào (gần 0?) rồi đối chiếu.

4Profiling

• Bộ nhớ đỉnh và thời gian một bước

  Đo với cùng mô hình, batch và độ dài, bật và tắt checkpointing, ghi thời gian bước và bộ nhớ đỉnh; tách phần thời gian tính lại trong backward. Chỉ kết luận ‘đáng đổi’ khi bộ nhớ tiết kiệm cho phép batch hoặc độ dài lớn hơn hoặc tránh các cách đắt hơn, đồng thời thông lượng token mỗi giây vẫn chấp nhận được.

• Tương đương gradient

  So gradient của đường tính lại với đường không tính lại trong một batch cố định; kiểm tra sai số tương đối nhỏ và nhất quán, đặc biệt ở đoạn có yếu tố ngẫu nhiên. Tính lại trùng khớp không có nghĩa bộ nhớ đã giảm; hai việc đo phải làm riêng.

• Độ hạt lựa chọn

  Thử vài mức độ hạt (mọi lớp, một nửa số lớp, chỉ attention) và vẽ bộ nhớ đỉnh với thời gian bước. Chọn điểm nằm trên đường biên: giảm thêm bộ nhớ sẽ phải trả bằng thời gian nhiều hơn rõ rệt.

Mẹo đo lường: ghim số liệu trên workload thật, ghi cả p50/p95/p99 thay vì chỉ trung bình, và luôn ghi lại điều kiện chạy (batch, cache, đồng thời, phần cứng) trước và sau khi tối ưu. Số đo thiếu điều kiện chạy thì không so sánh được.

5Ví dụ

6Thuật ngữ

Activation

Giá trị trung gian sinh ra khi chạy forward; backward cần chúng để tính gradient.

Recomputation (tính lại)

Chạy lại một đoạn forward lúc backward cần giá trị, thay vì giữ giá trị đó trong bộ nhớ suốt thời gian chờ.

Checkpoint boundary

Tensor ở ranh giới đoạn (thường là đầu vào mỗi lớp) được giữ lại làm điểm xuất phát để chạy lại đoạn đó.

Activation checkpointing

Tên kỹ thuật này (còn gọi gradient checkpointing); khác hẳn checkpoint là tệp lưu trọng số ở ô 3E.

Trạng thái ngẫu nhiên (RNG state)

Nếu đoạn tính lại có yếu tố ngẫu nhiên như dropout, phải khôi phục đúng trạng thái ngẫu nhiên để kết quả giống lần forward đầu.

7Bài liên quan

8Nguồn