Paralelisme data terpecah-pecah hibrida

Paralelisme data tersharded adalah teknik pelatihan terdistribusi hemat memori yang membagi keadaan 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, pustaka 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 Near-linear penskalaan blog pelatihan model raksasa untuk FSDP. AWS PyTorch

Pustaka SMP membuatnya mudah digunakan HYBRID_SHARD dan _HYBRID_SHARD_ZERO2 di seluruh jumlah GPU yang dapat dikonfigurasi, memperluas PyTorch FSDP asli yang mendukung sharding di satu node (HYBRID_SHARD) atau semua GPU (). FULL_SHARD PyTorch Panggilan FSDP dapat tetap apa adanya, dan Anda hanya perlu menambahkan hybrid_shard_degree argumen ke konfigurasi SMP, seperti yang ditunjukkan dalam contoh kode berikut. Anda tidak perlu mengubah nilai sharding_strategy argumen di pembungkus PyTorch FSDP di sekitar model Anda. PyTorch Anda dapat mener ShardingStrategy.HYBRID_SHARD uskan sebagai nilai. 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 hybrid_shard_degree parameter.

Cuplikan kode berikut menunjukkan cara menambahkan modul torch.sagemaker.init() inisialisasi SMP ke skrip pelatihan Anda dan mengatur kamus konfigurasi SMP dalam format JSON untuk peluncur pekerjaan pelatihan sambil mengikuti proses dua langkah yang diperkenalkan di. Gunakan perpustakaan paralelisme SageMaker 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 skrip 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 = ...