Đào tạo đa GPU

LeRobot huấn luyện trên nhiều GPU thông qua Hugging Face Accelerate. Ba bố cục song song dữ liệu được hỗ trợ:

Cài đặt

accelerate được bao gồm trong phần bổ sung training:

pip install 'lerobot[training]'

Ra mắt

Huấn luyện phân tán có thể được khởi chạy thông qua cả torchrun và accelerate launch. Accelerate được sử dụng như một trình khởi chạy đơn thuần: nó không quản lý cấu hình huấn luyện, và mọi cài đặt huấn luyện phân tán đều nằm trong hệ thống cấu hình riêng của LeRobot.

Với torchrun:

torchrun --nproc-per-node=2 $(which lerobot-train) \
  --dataset.repo_id=${HF_USER}/my_dataset \
  --policy.type=act \
  --policy.repo_id=${HF_USER}/my_trained_policy \
  --output_dir=outputs/train/act_multi_gpu \
  --job_name=act_multi_gpu \
  --trackio.enable=true

Với accelerate launch (dưới dạng trình khởi chạy đơn giản):

accelerate launch --num_processes=2 $(which lerobot-train) \
  --dataset.repo_id=${HF_USER}/my_dataset \
  --policy.type=act \
  --policy.repo_id=${HF_USER}/my_trained_policy \
  --output_dir=outputs/train/act_multi_gpu \
  --job_name=act_multi_gpu \
  --trackio.enable=true

Khi không có cờ --parallelism.*, lệnh khởi chạy đa tiến trình sẽ chạy DDP thông thường. Các lần chạy đa nút sử dụng các cờ chuẩn torchrun --nnodes/--node-rank/--rdzv-endpoint (hoặc accelerate launch --num_machines/--machine_rank/--main_process_ip).

Các tệp cấu hình YAML của Accelerate (accelerate launch --config_file some.yaml, accelerate config) không được hỗ trợ. Chúng cấu hình engine thông qua các biến môi trường, bỏ qua hệ thống cấu hình của LeRobot, do đó train_config.json sẽ không còn mô tả các cài đặt mà một lần chạy thực sự sử dụng. Vì vậy, lerobot-train sẽ từ chối khởi động khi các biến môi trường của accelerate được đặt. Thay vào đó, hãy đặt các cài đặt trong các cờ --parallelism.* / --accelerator.*, hoặc đặt LEROBOT_ALLOW_ACCELERATE_ENV=1 để xác nhận việc ghi đè và tiếp tục thực thi.

Ngữ nghĩa batch, tốc độ học và số bước

Mỗi worker song song dữ liệu trong dp_replicate × dp_shard nạp micro-batch --batch_size riêng của nó ở mỗi bước, do đó một bước huấn luyện tiêu thụ batch_size × dp_world_size mẫu, và × gradient_accumulation_steps trong số đó được đưa vào mỗi lần cập nhật optimizer:

effective_batch_size = batch_size × dp_world_size × gradient_accumulation_steps

Banner huấn luyện in ra phép phân tích này khi khởi động. --steps đếm số bước vòng lặp (micro-batch trên mỗi worker), không phải số lần cập nhật optimizer.

Gradient accumulation là một cờ cấp một:

torchrun --nproc-per-node=2 $(which lerobot-train) \
  --batch_size=8 --accelerator.gradient_accumulation.steps=4 ...

LeRobot không tự động điều chỉnh tốc độ học hoặc số bước khi kích thước batch hiệu dụng tăng lên. Nếu bạn mở rộng quy mô và muốn quá trình huấn luyện tương đương, vui lòng điều chỉnh thủ công, ví dụ với 2 GPU: nhân đôi --optimizer.lr (thang đo tuyến tính), hoặc giảm một nửa --steps.

Huấn luyện phân mảnh (FSDP)

Nếu một model quá lớn để huấn luyện với DDP, hãy shard nó bằng FSDP2:

torchrun --nproc-per-node=4 $(which lerobot-train) \
  --dataset.repo_id=${HF_USER}/my_dataset \
  --policy.type=<your_policy> \
  --parallelism.dp_shard=4 \
  --accelerator.mixed_precision=bf16 \
  --output_dir=outputs/train/my_policy_fsdp

--parallelism.dp_shard=-1 chia nhỏ (shard) trên số tiến trình mà trình khởi chạy đã khởi tạo.

--accelerator.mixed_precision hỗ trợ cả bf16 và fp16. Chi tiết trong mixed precision

Bọc các đơn vị

FSDP chia nhỏ mô hình theo các đơn vị (thường là khối transformer lặp lại) và thu thập từng đơn vị một trong quá trình forward/backward. Các policy khai báo đơn vị wrap của chúng thông qua _fsdp_wrap_modules trên lớp policy. Ví dụ, ACT khai báo ["ACTEncoderLayer", "ACTDecoderLayer"] và FastWAM khai báo ["MoTLayer"]. Đối với một policy không có khai báo _fsdp_wrap_modules, hãy truyền một trong các cờ bên dưới. Bạn có thể chỉ định tên lớp module một cách rõ ràng, hoặc sử dụng một policy dựa trên kích thước thay thế:

--accelerator.fsdp.wrap_modules='["MyTransformerBlock"]'   # explicit class names
--accelerator.fsdp.min_num_params=1000000                  # or: wrap every submodule above 1M params

FSDP không thể tách các trọng số liên kết (ví dụ: đầu ngôn ngữ được liên kết với nhúng đầu vào) giữa các đơn vị bọc. torch 2.12 và các phiên bản mới hơn sẽ dừng lại ở bước huấn luyện đầu tiên và nêu tên tham số, trong khi các phiên bản cũ hơn sẽ gỡ liên kết chúng một cách im lặng. Việc bọc lớp khối lặp lại thường giữ chúng cùng nhau trong đơn vị gốc.

Nếu một policy không khai báo _fsdp_wrap_modules và không có --accelerator.fsdp.wrap_modules hoặc --accelerator.fsdp.min_num_params được truyền vào, quá trình chạy sẽ thất bại ngay khi khởi động thay vì âm thầm chỉ bọc module gốc (điều này sẽ làm mất đi toàn bộ lợi ích tiết kiệm bộ nhớ do sharding).

Các cài đặt chia nhỏ khác:

• --accelerator.fsdp.reshard_after_forward: có nên giữ tham số của mỗi đơn vị ở bộ nhớ sau bước forward hay không.
• --accelerator.fsdp.cpu_offload: giữ các tham số, gradient và trạng thái optimizer trên CPU.
• --accelerator.fsdp.ignored_modules: một regex của các đường dẫn module cần giữ không bị chia nhỏ (unsharded).
• --accelerator.fsdp.mixed_precision_policy.{param_dtype,reduce_dtype,output_dtype,cast_forward_inputs}: các trường của MixedPrecisionPolicy trong PyTorch, được áp dụng cho mọi đơn vị FSDP. Khi bất kỳ trường nào được đặt, chính sách sẽ được xây dựng chỉ từ các trường này và thay thế chính sách được suy ra từ --accelerator.mixed_precision, do đó hãy đặt cả param_dtype. Trước torch 2.15, reduce_dtype không được đặt sẽ theo param_dtype. Từ torch 2.15, trường này sẽ theo dtype riêng của các tham số (float32 dưới mixed precision). float16 yêu cầu --accelerator.mixed_precision=fp16, vốn xây dựng loss scaler.

HSDP

Hybrid Sharded Data Parallel: các tham số, gradient và trạng thái optimizer được chia nhỏ (sharded) trên dp_shard ranks, và việc chia nhỏ đó được nhân bản (replicated) dp_replicate lần. Các thao tác all-gather tham số và reduce-scatter gradient diễn ra bên trong một nhóm shard; chỉ có thao tác all-reduce đồng bộ hóa các bản nhân bản mới diễn ra giữa các nhóm. Hai bậc tự do này phải nhân với nhau bằng kích thước world:

# 16 GPUs = 2 nodes × 8: shard within each node, replicate across nodes
torchrun --nnodes=2 --nproc-per-node=8 ... $(which lerobot-train) \
  --parallelism.dp_replicate=2 --parallelism.dp_shard=8 ...

PEFT

Chỉnh huấn luyện PEFT (--peft.*) hoạt động dưới FSDP2 và HSDP. Mô hình cơ sở bị đóng băng được chia mảnh cùng với các adapter, do đó mỗi rank chỉ giữ mảnh của trọng số cơ sở:

torchrun --nproc-per-node=4 $(which lerobot-train) \
  --policy.path=${HF_USER}/my_base_policy \
  --peft.method_type=LORA --peft.r=16 \
  --parallelism.dp_shard=4 \
  --accelerator.mixed_precision=bf16 \
  --output_dir=outputs/train/my_policy_fsdp_lora

• Các đơn vị được bọc vẫn được lấy từ khai báo _fsdp_wrap_modules của chính sách (hoặc --accelerator.fsdp.wrap_modules).
• Sử dụng mixed precision khi mô hình cơ sở được lưu trữ ở bf16 hoặc fp16. PEFT giữ các adapter của nó ở fp32 trong khi các bản sao của --peft.full_training_modules giữ dtype của mô hình cơ sở, và FSDP2 cần một dtype duy nhất cho các tham số có thể huấn luyện mà nó chia nhỏ cùng nhau. Khi bật mixed precision, accelerate sẽ upcast mọi tham số có thể huấn luyện lên fp32.
• Checkpoints lưu trữ adapter thay vì toàn bộ trọng số, trong mọi --checkpoint_format (xem Checkpoints).
• --accelerator.fsdp.ignored_modules được khớp với tên module của mô hình PEFT, bắt đầu bằng base_model.model.. Bắt đầu mẫu bằng .* để khớp với chúng.

Điểm kiểm tra

Mỗi checkpoint chứa một thư mục pretrained_model/ và một thư mục training_state/:

005000/  # the training step at that checkpoint
├── pretrained_model/
│   ├── config.json  # policy config
│   ├── train_config.json  # the full training config
│   ├── model.safetensors  # full weights (checkpoint_format ∈ {safetensors, safetensors_dcp}, or any non-sharded run)
│   ├── pytorch_model_fsdp_0/  # DCP weight shards (checkpoint_format ∈ {dcp, safetensors_dcp})
│   ├── policy_preprocessor.json  # preprocessor config (when the run has a preprocessor)
│   ├── policy_preprocessor_step_*.safetensors  # state of the stateful preprocessor steps
│   ├── policy_postprocessor.json  # postprocessor config (when the run has a postprocessor)
│   └── policy_postprocessor_step_*.safetensors  # state of the stateful postprocessor steps
└── training_state/
    ├── optimizer_param_groups.json  # optimizer param groups (non-sharded runs)
    ├── optimizer_state.safetensors  # optimizer state (non-sharded runs)
    ├── optimizer_0/  # DCP optimizer shards (sharded runs)
    ├── rng_state.safetensors  # rng states
    ├── scheduler_state.json  # scheduler state (when the run has a scheduler)
    ├── scaler_state.json  # fp16 loss-scaler state (mixed_precision=fp16 only)
    └── training_step.json  # step + dp_world_size/batch_size/grad_accum/precision + topology

Trong quá trình huấn luyện trên một GPU đơn hoặc DDP, pipeline sẽ serialize mỗi state dict vào một tệp duy nhất: model.safetensors cho model và optimizer_state.safetensors cho optimizer.

Trong quá trình huấn luyện phân mảnh, trạng thái của bộ tối ưu được lưu dưới dạng các phân mảnh DCP trong training_state/optimizer_0/, và bố cục của mô hình trong pretrained_model/ có thể được cấu hình thông qua --checkpoint_format:

Một lần chạy PEFT lưu adapter thay vì toàn bộ trọng số: adapter_model.safetensors và adapter_config.json thay thế cho model.safetensors, và các shard DCP chỉ chứa adapter. Mô hình cơ sở được nạp từ đường dẫn được ghi trong adapter_config.json.

Hai điều cần biết về các checkpoint (safetensors) đã được gom lại từ các lần chạy phân mảnh:

• Chúng lưu các trọng số fp32. Trong quá trình huấn luyện chính độ hỗn hợp, FSDP giữ một bản sao master fp32, và checkpoint lưu bản sao master này để đảm bảo việc tiếp tục huấn luyện diễn ra nhất quán.
• Gather là tập thể (tất cả các rank đều tham gia) nhưng chỉ tiến trình chính mới ghi.

Chuyển đổi checkpoint DCP

lerobot-convert-dcp gộp một thư mục shard DCP thành một model.safetensors thông thường, chạy ngoại tuyến và không cần GPU:

lerobot-convert-dcp --checkpoint_dir=outputs/train/run/checkpoints/005000
lerobot-convert-dcp --checkpoint_dir=... --delete_dcp=true --push_to_hub=${HF_USER}/my_policy

--push_to_hub công bố thư mục đã chuyển đổi dưới dạng kho lưu trữ mô hình. Các shard của một lần chạy PEFT được hợp nhất vào adapter_model.safetensors. Giống như bất kỳ checkpoint dcp nào, một checkpoint PEFT được lưu với --checkpoint_format=dcp chỉ có thể được nạp bằng --policy.path sau khi chuyển đổi này.

Tiếp tục

Tiếp tục với --resume=true --config_path=.../checkpoints/last/pretrained_model/train_config.json. Việc tiếp tục từ checkpoint DCP hỗ trợ resharding mô hình và trạng thái optimizer sang topology hiện tại, nghĩa là bạn có thể tiếp tục với cách chia dp_replicate/dp_shard khác. Data sampler luôn có thể tiếp tục ở đúng epoch và offset, nhưng chỉ chính xác theo mẫu khi world size và batch size khớp với lần chạy ban đầu (một cảnh báo sẽ được ghi log nếu không).

FSDP checkpoints được viết bởi LeRobot 0.6.x và các phiên bản trước đó sử dụng bố cục trên đĩa khác (trạng thái optimizer đầy đủ đã được gom lại) và không thể được tiếp tục.

Ghi chú

• Việc lưu checkpoint và xuất bản khi kết thúc huấn luyện là các thao tác tập thể (mọi rank đều tham gia). Các trọng số đã gom lại, các tệp phụ và các bản tải lên Hub được viết riêng bởi tiến trình chính. Việc lưu checkpoint và xuất bản khi kết thúc huấn luyện là các thao tác tập thể (mọi rank đều tham gia). Các trọng số đã gom lại, các tệp phụ và các bản tải lên Hub được viết riêng bởi tiến trình chính.
• Các chỉ số được giảm (reduce) trên các rank trước khi ghi log: các giá trị loss được lấy trung bình, và samples/s báo cáo thông lượng trên toàn cụm.
• Việc lập lịch learning-rate được thực hiện một bước cho mỗi bước huấn luyện, bất kể số lượng tiến trình (step_scheduler_with_optimizer=False được tích hợp sẵn).

Để tìm hiểu về cơ chế nền tảng, hãy xem hướng dẫn FSDP của Accelerate. Để đi sâu hơn vào huấn luyện quy mô lớn, hãy tham khảo Ultrascale Playbook.