Các phương pháp huấn luyện độ chính xác thấp

Accelerate cung cấp các tích hợp để huấn luyện bằng các phương pháp độ chính xác thấp hơn trên phần cứng được hỗ trợ thông qua các gói TransformersEngine và torchao. Tài liệu này sẽ hướng dẫn bạn phần cứng nào được hỗ trợ, cách cấu hình Accelerator của bạn để tận dụng các phương pháp độ chính xác thấp, và những gì bạn có thể kỳ vọng khi huấn luyện.

Ý nghĩa của việc huấn luyện trên FP8

Để tìm hiểu sâu hơn về các chi tiết trong quá trình huấn luyện với FP8 bằng PyTorch và Accelerate, hãy xem concept_guide về lý do tại sao điều này có thể gây khó khăn. Về cơ bản, thay vì huấn luyện với BF16, một số (hoặc toàn bộ) khía cạnh của việc huấn luyện mô hình có thể được thực hiện bằng 8 bit thay vì 16 bit. Thách thức là làm được điều đó mà không làm giảm hiệu suất cuối cùng.

Tính năng này chỉ được bật trên các phần cứng NVIDIA cụ thể, bao gồm:

• Bất kỳ card đồ họa tiêu dùng nào thuộc dòng 3000 trở lên (chẳng hạn như 4090)
• Các kiến trúc GPU dựa trên Hopper (chẳng hạn như H100 và H200)

Điều này sẽ dẫn đến việc giảm một phần bộ nhớ được sử dụng (vì chúng ta đã cắt giảm một nửa bộ nhớ cần thiết cho một số phần của quá trình huấn luyện) và tăng thông lượng should cũng sẽ được nhìn thấy đối với các mô hình lớn hơn có thể thay thế một số lớp bằng các lớp hỗ trợ FP8.

Cấu hình Accelerator

Hiện tại, hai backend cho FP8 đang được bảo trì tích cực được hỗ trợ (TransformersEngine và torchao), mỗi backend có các khả năng và cấu hình khác nhau.

Để sử dụng một trong hai, cùng một API cốt lõi được dùng. Chỉ cần truyền mixed_precision="fp8" vào Accelerator, trong quá trình accelerate config khi được hỏi về độ chính xác hỗn hợp, hoặc như một phần của tệp config.yaml của bạn ở khóa mixed_precision:

from accelerate import Accelerator
accelerator = Accelerator(mixed_precision="fp8")

Để chỉ định một backend (và tùy chỉnh các phần khác của cấu hình độ chính xác hỗn hợp FP8), bạn có thể sử dụng một trong các dataclass RecipeKwargs như utils.AORecipeKwargs hoặc utils.TERecipeKwargs; bạn cũng có thể làm rõ điều này trong file cấu hình yaml/trong quá trình accelerate launch. Chúng tôi khuyến nghị sử dụng TransformersEngine hoặc torchao cho các dự án mới:

from accelerate import Accelerator
from accelerate.utils import TERecipeKwargs, AORecipeKwargs
# Use TransformersEngine
kwargs = [TERecipeKwargs()]
# Or to use torchao
# kwargs = [AORecipeKwargs()]
accelerator = Accelerator(mixed_precision="fp8", kwarg_handlers=kwargs)

mixed_precision: fp8
fp8_config:
  amax_compute_algo: max
  amax_history_len: 1024
  backend: TE
  fp8_format: HYBRID
  interval: 1
  margin: 0
  override_linear_precision: (false, false, false)
  use_autocast_during_eval: false

Cấu hình TransformersEngine

TransformersEngine có nhiều tùy chọn để tùy chỉnh cách thức và nội dung của các phép tính FP8. Danh sách đầy đủ các đối số được hỗ trợ và ý nghĩa của chúng có sẵn trong tài liệu của NVIDIA, tuy nhiên chúng được nêu lại trong docstring của FP8KwargsHandler để bạn tiện tham khảo.

Accelerate cố gắng thiết lập các giá trị mặc định hợp lý, nhưng việc tự khám phá và tinh chỉnh các tham số khác nhau có thể giúp cải thiện hiệu suất.

Để sử dụng, hãy chỉ định backend="te" và sửa đổi bất kỳ đối số nào bạn muốn như một phần của bộ xử lý kwarg của bạn:

from accelerate import Accelerator
from accelerate.utils import TERecipeKwargs
kwargs = [TERecipeKwargs(...)]
accelerator = Accelerator(mixed_precision="fp8", kwargs_handlers=kwargs)

Hoặc trong quá trình accelerate launch thông qua --fp8_backend=te .... Sử dụng accelerate launch --fp8_backend=te -h để xem các đối số liên quan.

Tương tự, bạn có thể thiết lập điều này trong config.yaml của mình:

mixed_precision: fp8
fp8_config:
    amax_compute_algo: max
    amax_history_len: 1024
    backend: TE
    fp8_format: HYBRID
    interval: 1
    margin: 0
    override_linear_precision: (false, false, false)
    use_autocast_during_eval: false

Cấu hình torchao

torchao là một backend FP8 có thể tùy chỉnh, được xây dựng dựa trên PyTorch, với mục tiêu dễ tiếp cận hơn so với hai engine trước đó. Một trong những khác biệt cốt lõi của ao so với hai engine trước là, để đảm bảo ổn định về mặt số học, việc giữ nguyên độ chính xác thông thường (dù là FP32 hay BF16) cho lớp đầu tiên và lớp cuối cùng của mô hình, trong khi lượng tử hóa các lớp còn lại xuống FP8, thường cho kết quả tốt hơn. Do đó, cấu hình cho ao có phần khác biệt:

Lưu ý: API này mang tính thử nghiệm và có thể thay đổi

from accelerate import Accelerator
from accelerate.utils import AORecipeKwargs, TorchDynamoPlugin, FullyShardedDataParallelPlugin
from torchao.float8 import Float8LinearConfig

fsdp2_plugin = FullyShardedDataParallelPlugin(
  fsdp_version=2,
  cpu_ram_efficient_loading=False, # CPU RAM efficient loading CANNOT work with fp8 torchao
  fsdp_auto_wrap_policy="TRANSFORMER_BASED_WRAP",
)
dynamo_plugin = TorchDynamoPlugin(
  backend="inductor",
  use_regional_compilation=True,
)
fp8_config = Float8LinearConfig(
  enable_fsdp_float8_all_gather=True, # Use FP8 all_gather in FSDP2
  pad_inner_dim=True,
)
kwargs = [AORecipeKwargs(
  config=fp8_config
)]
accelerator = Accelerator(
  mixed_precision="fp8",
  fsdp_plugin=fsdp2_plugin,
  dynamo_plugin=dynamo_plugin,
  kwarg_handlers=kwargs,
)

Hoặc trong quá trình accelerate launch thông qua --fp8_backend=ao .... Sử dụng accelerate launch --fp8_backend=ao -h để xem các đối số liên quan.

Tương tự, điều này có thể được thiết lập trong config.yaml:

mixed_precision: fp8
fsdp_config:
  fsdp_auto_wrap_policy: TRANSFORMER_BASED_WRAP
  fsdp_cpu_ram_efficient_loading: false
  fsdp_version: 2
fp8_config:
  backend: AO
  pad_inner_dim: true
  enable_fsdp_float8_all_gather: true
dynamo_config:
  dynamo_backend: INDUCTOR
  dynamo_use_regional_compilation: true

Để tìm hiểu thêm về các tham số cụ thể cần sử dụng, vui lòng xem kho torchao chính thức.

Ví dụ Zoo

Chúng tôi có các ví dụ minh họa việc huấn luyện với FP8, cả với accelerate lẫn cài đặt cơ bản của nó, có sẵn trong kho accelerate. Hiện tại chúng tôi hỗ trợ các script minh họa:

• Một GPU
• Song song dữ liệu phân tán (Đa GPU)
• Song song dữ liệu phân mảnh hoàn toàn
• DeepSpeed ZeRO 1 đến 3

Tìm hiểu thêm tại đây

Đọc thêm

Để tìm hiểu thêm về huấn luyện trong FP8, hãy xem các tài nguyên sau:

• Hướng dẫn khái niệm của chúng tôi đi sâu hơn về TransformersEngine và torchao
• Tài liệu của transformers-engine
• Tài liệu của torchao