Paralelismo híbrido de dados fragmentados

O paralelismo de dados fragmentados é uma técnica de treinamento distribuído que economiza memória e divide o estado de um modelo (parâmetros do modelo, gradientes e estados do otimizador) entre dispositivos. Isso ajuda você a ajustar um modelo maior ou aumentar o tamanho do lote usando a memória liberada. GPU A SMP biblioteca oferece a capacidade de executar paralelismo de dados fragmentados com o PyTorch Fully Sharded Data Parallel (). FSDP PyTorch FSDPpor padrão, fragmentos em todo o conjunto de fragmentos que GPUs estão sendo usados. Na SMP v2, a biblioteca oferece esse paralelismo de dados fragmentados, além de estender a fragmentação PyTorch híbrida (HYBRID_SHARD), que é uma das estratégias de fragmentação fornecidas PyTorch FSDP por:,,,. PyTorch FSDP FULL_SHARD SHARD_GRAD_OP HYBRID_SHARD _HYBRID_SHARD_ZERO2 Estender a fragmentação híbrida dessa maneira ajuda a implementar, scale-aware-sharding conforme descrito no blog Escalonamento quase linear do treinamento de modelos gigantes em AWSpara PyTorchFSDP.

A SMP biblioteca facilita o uso HYBRID_SHARD e _HYBRID_SHARD_ZERO2 abrange qualquer número configurável deGPUs, estendendo o nativo PyTorch FSDP que suporta fragmentação em um único nó (HYBRID_SHARD) ou em todos GPUs (). FULL_SHARD PyTorch FSDPas chamadas podem permanecer como estão, e você só precisa adicionar o hybrid_shard_degree argumento à SMP configuração, conforme mostrado no exemplo de código a seguir. Você não precisa alterar o valor do sharding_strategy argumento no PyTorch FSDP invólucro ao redor do seu PyTorch modelo. Você pode passar ShardingStrategy.HYBRID_SHARD como valor. Como alternativa, a SMP biblioteca substitui a estratégia no script e a define como ShardingStrategy.HYBRID_SHARD se você especificar um valor igual ou maior que 2 para o hybrid_shard_degree parâmetro.

Os trechos de código a seguir mostram como adicionar o módulo de SMP inicialização torch.sagemaker.init() ao seu script de treinamento e configurar o dicionário de SMP configuração no JSON formato para o lançador de tarefas de treinamento, seguindo o processo de duas etapas apresentado em. Comece com a biblioteca de paralelismo de SageMaker modelos v2 Você não precisa fazer nenhuma alteração em seu PyTorch modelo ou PyTorch FSDPconfiguração. Para obter mais informações sobre o parâmetro hybrid_shard_degree, consulte SMPparâmetros de configuração do recurso principal v2.

SMPdicionário de configuração

{ "hybrid_shard_degree": 16 }

No roteiro de treinamento

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 = ...