Paralelisme data sharded hibrida

Paralelisme data sharded adalah teknik pelatihan terdistribusi hemat memori yang membagi status model (parameter model, gradien, dan status pengoptimal) di seluruh perangkat. Ini membantu Anda menyesuaikan model yang lebih besar atau meningkatkan ukuran batch menggunakan memori GPU yang dibebaskan. Pustaka SMP menawarkan kemampuan menjalankan paralelisme data sharded dengan PyTorch Fully Sharded Data Parallel (FSDP). PyTorch FSDP secara default pecahan di seluruh rangkaian GPU yang digunakan. Di SMP v2, perpustakaan menawarkan paralelisme data sharded ini di atas PyTorch FSDP dengan memperluas sharding PyTorch hybrid (HYBRID_SHARD), yang merupakan salah satu strategi sharding yang disediakan oleh FSDP:,,,. PyTorch FULL_SHARD SHARD_GRAD_OP HYBRID_SHARD _HYBRID_SHARD_ZERO2 Memperluas sharding hibrida dengan cara ini membantu mengimplementasikan scale-aware-sharding seperti yang dijelaskan dalam blog Penskalaan linier dekat pelatihan model raksasa untuk FSDP. AWS PyTorch

Pustaka SMP membuatnya mudah digunakan HYBRID_SHARD dan _HYBRID_SHARD_ZERO2 di sejumlah GPU yang dapat dikonfigurasi, memperluas PyTorch FSDP asli yang mendukung sharding di satu node () atau semua GPU ()HYBRID_SHARD. FULL_SHARD PyTorch Panggilan FSDP dapat tetap apa adanya, dan Anda hanya perlu menambahkan hybrid_shard_degree argumen ke konfigurasi SMP, seperti yang ditunjukkan pada contoh kode berikut. Anda tidak perlu mengubah nilai sharding_strategy argumen dalam pembungkus PyTorch FSDP di sekitar model Anda. PyTorch Anda dapat lulus ShardingStrategy.HYBRID_SHARD sebagai nilainya. Atau, pustaka SMP mengganti strategi dalam skrip dan menyetelnya ShardingStrategy.HYBRID_SHARD jika Anda menentukan nilai yang sama dengan atau lebih besar dari 2 ke parameter. hybrid_shard_degree

Cuplikan kode berikut menunjukkan cara menambahkan modul inisialisasi SMP torch.sagemaker.init() ke skrip pelatihan Anda dan mengatur kamus konfigurasi SMP dalam format JSON untuk peluncur pekerjaan pelatihan sambil mengikuti proses dua langkah yang diperkenalkan. Memulai dengan SageMaker pustaka paralelisme model v2 Anda tidak perlu membuat perubahan apa pun pada PyTorch model atau konfigurasi PyTorch FSDP Anda. Untuk informasi tentang parameter hybrid_shard_degree, lihat Parameter konfigurasi fitur inti SMP v2.

Kamus konfigurasi SMP

{ "hybrid_shard_degree": 16 }

Dalam naskah pelatihan

import torch.sagemaker as tsm
tsm.init()

# Set up a PyTorch model
model = ...

# Wrap the PyTorch model using the PyTorch FSDP module
model = FSDP(
    model,
    ...
)

# Optimizer needs to be created after FSDP wrapper
optimizer = ...