As traduções são geradas por tradução automática. Em caso de conflito entre o conteúdo da tradução e da versão original em inglês, a versão em inglês prevalecerá.

Apontando pontos de verificação e ajustando um modelo com paralelismo de modelos

A biblioteca de paralelismo de SageMaker modelos fornece APIs de ponto de verificação para salvar o estado do modelo e o estado do otimizador divididos pelas várias estratégias de paralelismo do modelo e para carregar pontos de verificação para treinamento contínuo de onde você deseja reiniciar o treinamento e ajustar. As APIs também oferecem opções de suporte para salvar parcialmente ou totalmente os estados do modelo e do otimizador.

Pontos de verificação de um modelo distribuído

Escolha um dos tópicos a seguir, dependendo da estrutura entre PyTorch e TensorFlow e da versão da biblioteca de paralelismo de SageMaker modelos que você usa.

Apontando um PyTorch modelo distribuído (para a biblioteca de paralelismo de SageMaker modelos v1.10.0 e posterior)

A biblioteca de paralelismo de SageMaker modelos fornece APIs de ponto de verificação para salvar e carregar pontos de verificação completos ou parciais do estado do modelo distribuído e do estado do otimizador.

Pontos de verificação parciais

Para salvar pontos de verificação de um treinamento de modelos com paralelismo de modelos, use a API smdistributed.modelparallel.torch.save_checkpoint com a opção de ponto de verificação parcial definida como true (partial=True). Isto salva cada partição de modelos individualmente. Além do modelo e do estado do otimizador, você também pode salvar quaisquer dados personalizados adicionais por meio do argumento user_content. O modelo com ponto de verificação, o otimizador e o conteúdo do usuário são salvos como arquivos separados. A chamada de API save_checkpoint cria pastas de pontos de verificação na estrutura a seguir.

- 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)

Para retomar o treinamento a partir de pontos de verificação parciais, use a API smdistributed.modelparallel.torch.resume_from_checkpoint com partial=True e especifique o diretório do ponto de verificação e a tag usada ao salvar os pontos de verificação parciais. Observe que o carregamento real dos pesos do modelo ocorre após o particionamento do modelo, durante a primeira execução da step function de treinamento decorada smdistributed.modelparallel.torch.step.

Ao salvar um ponto de verificação parcial, a biblioteca também salva a decisão da partição do modelo como arquivos com extensão de arquivo .pt. Por outro lado, ao retomar o ponto de verificação parcial, a biblioteca carrega os arquivos de decisão de partição juntos. Depois que a decisão de partição é carregada, não é possível alterar a partição.

O trecho de código a seguir mostra como definir as APIs do ponto de verificação em um script de treinamento. 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
)

Pontos de verificação totais

Para salvar o artefato do modelo final para fins de inferência, use a API smdistributed.modelparallel.torch.save_checkpoint com partial=False, que combinam as partições do modelo para criar um único artefato do modelo. Observe que isso não combina os estados do otimizador.

Para inicializar o treinamento com pesos específicos, considerando um ponto de verificação completo do modelo, você pode usar a API smdistributed.modelparallel.torch.resume_from_checkpoint com partial=False. Observe que isso não combina os estados de carregamento do otimizador.

O código a seguir mostra um exemplo de como usar as APIs de ponto de verificação para verificar totalmente um PyTorch modelo treinado com paralelismo de modelos.

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
)

Apontando um PyTorch modelo distribuído (para a biblioteca de paralelismo de SageMaker modelos entre v1.6.0 e v1.9.0)

A biblioteca de paralelismo de SageMaker modelos fornece funções Python para salvar pontos de verificação parciais ou completos para treinar trabalhos com paralelismo de tensores. O procedimento a seguir mostra como usar o smp.save() e smp.load() para salvar e carregar um ponto de verificação ao usar o paralelismo de tensores.

Verificando um modelo distribuído TensorFlow

Para salvar um TensorFlow modelo durante o treinamento com o paralelismo de modelos, use as seguintes funções fornecidas pela biblioteca de paralelismo de SageMaker modelos.

Ajuste de um modelo distribuído

O ajuste fino precisa ser configurado em seu script de treinamento. O trecho de código a seguir mostra um exemplo de estrutura de um script de treinamento usando a classe AutoModelForCausalLM de Hugging Face Transformers com modificações para registrar os módulos e as configurações para ajuste fino. smdistributed.model.parallel.torch

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
    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)

Para obter um exemplo completo de scripts de treinamento e notebooks Jupyter, consulte os exemplos do GPT-2 no repositório Examples. PyTorch SageMaker GitHub