Terjemahan disediakan oleh mesin penerjemah. Jika konten terjemahan yang diberikan bertentangan dengan versi bahasa Inggris aslinya, utamakan versi bahasa Inggris.

Checkpointing dan Fine-Tuning Model dengan Paralelisme Model

Pustaka paralelisme SageMaker model menyediakan pos pemeriksaan APIs untuk menyimpan status model dan status pengoptimal yang dibagi oleh berbagai strategi paralelisme model, dan untuk memuat pos pemeriksaan untuk pelatihan berkelanjutan dari tempat Anda ingin memulai ulang pelatihan dan menyempurnakan. Opsi APIs juga mendukung untuk menyimpan status model dan pengoptimal sebagian atau seluruhnya.

Checkpointing model terdistribusi

Pilih salah satu topik berikut tergantung pada kerangka kerja antara PyTorch dan TensorFlow dan versi pustaka paralelisme SageMaker model yang Anda gunakan.

Checkpointing PyTorch model terdistribusi (untuk pustaka paralelisme SageMaker model v1.10.0 dan yang lebih baru)

Pustaka paralelisme SageMaker model menyediakan pos pemeriksaan APIs untuk menyimpan dan memuat pos pemeriksaan penuh atau sebagian dari status model terdistribusi dan status pengoptimalnya.

Pos pemeriksaan sebagian

Untuk menyimpan pos pemeriksaan model yang dilatih dengan paralelisme model, gunakan smdistributed.modelparallel.torch.save_checkpointAPI dengan opsi checkpointing sebagian yang disetel ke true (). partial=True Ini menghemat setiap partisi model satu per satu. Selain model dan status pengoptimal, Anda juga dapat menyimpan data kustom tambahan melalui user_content argumen. Model checkpoint, optimizer, dan konten pengguna disimpan sebagai file terpisah. Panggilan save_checkpoint API membuat folder pos pemeriksaan dalam struktur berikut.

- path
  - ${tag}_partial (folder for partial checkpoints)
    - model_rankinfo.pt
    - optimizer_rankinfo.pt
    - fp16_states_rankinfo.pt
    - user_content.pt
  - $tag (checkpoint file for full checkpoints)
  - user_content_$tag (user_content file for full checkpoints)
  - newest (a file that indicates the newest checkpoint)

Untuk melanjutkan pelatihan dari pos pemeriksaan sebagian, gunakan smdistributed.modelparallel.torch.resume_from_checkpointAPI denganpartial=True, dan tentukan direktori pos pemeriksaan dan tag yang digunakan saat menyimpan pos pemeriksaan sebagian. Perhatikan bahwa pemuatan bobot model yang sebenarnya terjadi setelah partisi model, selama menjalankan pertama fungsi langkah pelatihan smdistributed.modelparallel.torch.step yang didekorasi.

Saat menyimpan pos pemeriksaan sebagian, perpustakaan juga menyimpan keputusan partisi model sebagai file dengan ekstensi .pt file. Sebaliknya, ketika melanjutkan dari pos pemeriksaan sebagian, perpustakaan memuat file keputusan partisi bersama-sama. Setelah keputusan partisi dimuat, Anda tidak dapat mengubah partisi.

Cuplikan kode berikut menunjukkan cara mengatur pos pemeriksaan APIs dalam skrip pelatihan. PyTorch

import smdistributed.modelparallel.torch as smp

model = ...
model = smp.DistributedModel(model)
optimizer = ...
optimizer = smp.DistributedOptimizer(optimizer)
user_content = ...     # additional custom data
checkpoint_path = "/opt/ml/checkpoint/model_parallel"

# Save a checkpoint.
smp.save_checkpoint(
    path=checkpoint_path,
    tag=f"total_steps{total_steps}",
    partial=True,
    model=model,
    optimizer=optimizer,
    user_content=user_content
    num_kept_partial_checkpoints=5
)

# Load a checkpoint.
# This automatically loads the most recently saved checkpoint.
smp_checkpoint = smp.resume_from_checkpoint(
    path=checkpoint_path, 
    partial=True
)

Checkpointing penuh

Untuk menyimpan artefak model akhir untuk tujuan inferensi, gunakan smdistributed.modelparallel.torch.save_checkpoint API denganpartial=False, yang menggabungkan partisi model untuk membuat artefak model tunggal. Perhatikan bahwa ini tidak menggabungkan status pengoptimal.

Untuk menginisialisasi pelatihan dengan bobot tertentu, dengan pos pemeriksaan model lengkap, Anda dapat menggunakan API dengansmdistributed.modelparallel.torch.resume_from_checkpoint. partial=False Perhatikan bahwa ini tidak memuat status pengoptimal.

Kode berikut menunjukkan contoh cara menggunakan pos pemeriksaan APIs untuk memeriksa sepenuhnya model yang dilatih dengan PyTorch paralelisme model.

import smdistributed.modelparallel.torch as smp

model = ...
model = smp.DistributedModel(model)
optimizer = ...
optimizer = smp.DistributedOptimizer(optimizer)
user_content = ...     # additional custom data
checkpoint_path = "/opt/ml/checkpoint/model_parallel"

# Save a checkpoint.
smp.save_checkpoint(
    path=checkpoint_path,
    tag=f"total_steps{total_steps}",
    partial=False,
    model=model,
    optimizer=optimizer,
    user_content=user_content
    num_kept_partial_checkpoints=5
)

# Load a checkpoint.
# This automatically loads the most recently saved checkpoint.
smp_checkpoint = smp.resume_from_checkpoint(
    path=checkpoint_path, 
    partial=False
)

Checkpointing PyTorch model terdistribusi (untuk pustaka paralelisme SageMaker model antara v1.6.0 dan v1.9.0)

Pustaka paralelisme SageMaker model menyediakan fungsi Python untuk menyimpan pos pemeriksaan sebagian atau penuh untuk pekerjaan pelatihan dengan paralelisme tensor. Prosedur berikut menunjukkan cara menggunakan smp.save()dan menyimpan dan smp.load()memuat pos pemeriksaan saat Anda menggunakan paralelisme tensor.

Checkpointing model terdistribusi TensorFlow

Untuk menyimpan TensorFlow model saat berlatih dengan paralelisme model, gunakan fungsi berikut yang disediakan oleh perpustakaan paralelisme SageMaker model.

Menyetel model terdistribusi

Penyetelan halus perlu dikonfigurasi dalam skrip pelatihan Anda. Cuplikan kode berikut menunjukkan contoh struktur skrip pelatihan menggunakan kelas AutoModelForCausalLM dari Hugging Face Transformers dengan modifikasi untuk mendaftarkan smdistributed.model.parallel.torch modul dan pengaturan untuk fine-tuning.

import argparse
from transformers import AutoModelForCausalLM
import smdistributed.modelparallel
import smdistributed.modelparallel.torch as smp

def parse_args():

    parser = argparse.ArgumentParser()

    # set an arg group for model
    model_grp = parser.add_argument_group(
        title="model", description="arguments to describe model configuration"
    )

    ... # set up numerous args to parse from the configuration dictionary to the script for training

    # add arg for activating fine-tuning
    model_grp.add_argument(
        "--fine_tune",
        type=int,
        default=0,
        help="Fine-tune model from checkpoint or pretrained model",
    )

def main():
    """Main function to train GPT."""
    args = parse_args()

    ... # parse numerous args

    if args.fine_tune > 0 and args.delayed_param > 0 and smp.rank() == 0:
        pretrained_model = AutoModelForCausalLM.from_pretrained(
            args.model_name or args.model_dir
        )
        model_state_dict = pretrained_model.state_dict()
        path = os.path.join(args.model_dir, "fullmodel.pt")
        torch.save(model_state_dict, path)

    # create a Transformer model and wrap by smp.model_creation() 
    # with options to configure model parallelism parameters offered by SageMaker AI
    with smp.model_creation(
        tensor_parallelism=smp.tp_size() > 1 or args.use_distributed_transformer > 0,
        zero_init=args.use_distributed_transformer == 0,
        dtype=dtype,
        distribute_embedding=args.sharded_data_parallel_degree > 1 and smp.tp_size() > 1,
        use_alibi=args.alibi > 0,
        attention_in_fp32=args.attention_in_fp32 > 0,
        fp32_residual_addition=args.residual_addition_in_fp32 > 0,
        query_key_layer_scaling=args.query_key_layer_scaling > 0 and args.bf16 < 1,
        fused_softmax=args.fused_softmax > 0,
        fused_dropout=args.fused_dropout > 0,
        fused_bias_gelu=args.fused_bias_gelu > 0,
        flash_attention=args.flash_attention > 0,
    ):
        if args.fine_tune > 0 and args.delayed_param == 0:
            model = AutoModelForCausalLM.from_pretrained(
                args.model_name or args.model_dir
            )
        else:
            model = AutoModelForCausalLM.from_config(model_config)

    # wrap the model by smp.DistributedModel() to apply SageMaker model parallelism
    model = smp.DistributedModel(
        model, trace_device="gpu", backward_passes_per_step=args.gradient_accumulation
    )

    # wrap the optimizer by smp.DistributedOptimizer() to apply SageMaker model parallelism
    optimizer= ... # define an optimizer
    optimizer = smp.DistributedOptimizer(
        optimizer,
        static_loss_scale=None,
        dynamic_loss_scale=True,
        dynamic_loss_args={"scale_window": 1000, "min_scale": 1, "delayed_shift": 2},
    )

    # for fine-tuning, use smp.resume_from_checkpoint() to load a pre-trained model
    if args.fine_tune > 0 and args.delayed_param > 0:
        smp.resume_from_checkpoint(args.model_dir, tag="fullmodel.pt", partial=False)

Untuk contoh lengkap skrip pelatihan dan notebook Jupyter, lihat contoh GPT-2 di repositori Contoh AI. PyTorch SageMaker GitHub