Inisialisasi parameter tertunda

Inisialisasi model besar untuk pelatihan tidak selalu memungkinkan dengan memori GPU yang terbatas. Untuk mengatasi masalah memori GPU yang tidak mencukupi ini, Anda dapat menginisialisasi model pada memori CPU. Namun, untuk model yang lebih besar dengan lebih dari 20 atau 40 miliar parameter, bahkan memori CPU mungkin tidak cukup. Untuk kasus seperti itu, kami sarankan Anda menginisialisasi model pada apa yang PyTorch memang gil perangkat meta, yang memungkinkan pembuatan tensor tanpa data yang melekat padanya. Tensor pada perangkat meta hanya membutuhkan informasi bentuk, dan ini memungkinkan untuk membuat model besar dengan parameternya pada perangkat meta. Hugging Face Acceler ate menyediakan pengelola konteks init_empty_weights untuk membantu membuat model seperti itu pada perangkat meta sambil menginisialisasi buffer pada perangkat biasa. Sebelum pelatihan dimulai, PyTorch FSDP menginisialisasi parameter model. Fitur inisialisasi parameter tertunda dari SMP v2 ini menunda pembuatan parameter model ini terjadi setelah PyTorch FSDP melakukan sharding parameter. PyTorch FSDP menerima fungsi inisialisasi parameter (param_init_fn) saat sharding modul, dan memang param_init_fn gil setiap modul. param_init_fnAPI mengambil modul sebagai argumen dan menginisialisasi semua parameter di dalamnya, tidak termasuk parameter modul anak mana pun. Perhatikan bahwa perilaku ini berbeda dari PyTorch v2.0.1 asli yang memiliki bug yang menyebabkan parameter diinisialisasi beberapa kali.

SMP v2 menyediakan torch.sagemaker.delayed_param. DelayedParamIniter API untuk menerapkan inisialisasi parameter tertunda.

Cuplikan kode berikut menunjukkan cara menerapkan torch.sagemaker.delayed_param.DelayedParamIniter API ke skrip pelatihan Anda.

Asumsikan bahwa Anda memiliki skri PyTorch p pelatihan FSDP sebagai berikut.

# Creation of model on meta device
from accelerate import init_empty_weights
with init_empty_weights():
    model = create_model()

# Define a param init fn, below is an example for Hugging Face GPTNeoX.
def init_weights(module):
    d = torch.cuda.current_device()
    # Note that below doesn't work if you have buffers in the model
    # buffers will need to reinitialized after this call
    module.to_empty(device=d, recurse=False)
    if isinstance(module, (nn.Linear, Conv1D)):
        module.weight.data.normal_(mean=0.0, std=args.initializer_range)
        if module.bias:
            module.bias.data.zero_()
    elif isinstance(module, nn.Embedding):
        module.weight.data.normal_(mean=0.0, std=args.initializer_range)
        if module.padding_idx:
            module.weight.data[module.padding_idx].zero_()
    elif isinstance(module, nn.LayerNorm):
        module.bias.data.zero_()
        module.weight.data.fill_(1.0)

# Changes to FSDP wrapper.
model = FSDP(
    model,
    ...,
    param_init_fn=init_weights
)

# At this point model is initialized and sharded for sharded data parallelism.

Perhatikan bahwa pendekatan inisialisasi parameter tertunda bukanlah model agnostik. Untuk mengatasi masalah ini, Anda perlu menulis init_weights fungsi seperti yang ditunjukkan pada contoh sebelumnya untuk mencocokkan inisialisasi dalam definisi model asli, dan itu harus mencakup semua parameter model. Untuk menyederhanakan proses persiapan init_weights fungsi tersebut, SMP v2 mengimplementasikan fungsi inisialisasi ini untuk model berikut: GPT-2,, GPT-J GPT-NeoX, dan Llama dari Hugging Face Transformers. torch.sagemaker.delayed_param.DelayedParamIniterAPI juga bekerja dengan implementasi paralel tensor SMP, torch.sagemaker.tensor_parallel.transformer.TransformerLMHead model, yang dapat Anda panggil setelah panggilan obor.sagemaker.transform API.

Dengan menggunakan torch.sagemaker.delayed_param.DelayedParamIniter API, Anda dapat menyesuaikan skri PyTorch p FSDP Anda sebagai berikut. Setelah membuat model dengan bobot kosong, daftar torch.sagemaker.delayed_param.DelayedParamIniter kan API ke model, dan tentukan objeknya. Meneruskan objek ke param_init_fn kelas PyTorch FSDP.

from torch.sagemaker.delayed_param import DelayedParamIniter
from accelerate import init_empty_weights

with init_empty_weights():
    model = create_model()
    
delayed_initer = DelayedParamIniter(model)

with delayed_initer.validate_params_and_buffers_inited():
    model = FSDP(
        model,
        ...,
        param_init_fn=delayed_initer.get_param_init_fn()
    )

Catatan tentang bobot yang diikat

Saat melatih model dengan beban terikat, kita perlu berhati-hati untuk mengikat bobot setelah menginisialisasi bobot dengan inisialisasi parameter yang tertunda. PyTorch FSDP tidak memiliki mekanisme untuk mengikat bobot setelah menginisialisasi menggunakan param_init_fn seperti di atas. Untuk mengatasi kasus seperti itu, kami menambahkan API untuk mengiz post_init_hook_fn inkan a, yang dapat digunakan untuk mengikat bobot. Anda dapat meneruskan fungsi apa pun di sana yang menerima modul sebagai argumen, tetapi kami juga memiliki definisi yang telah post_param_init_fn ditentukan sebelumnya di DelayedParamIniter mana memanggil tie_weights metode modul jika ada. Perhatikan bahwa aman untuk selalu masuk post_param_init_fn meskipun tidak ada tie_weights metode untuk modul.

with delayed_initer.validate_params_and_buffers_inited():
    model = FSDP(
        model,
        ...,
        param_init_fn=delayed_initer.get_param_init_fn(),
        post_param_init_fn=delayed_initer.get_post_param_init_fn()
    )