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 da GPU. A biblioteca SMP oferece a capacidade de executar paralelismo de dados fragmentados com o PyTorch Fully Sharded Data Parallel (FSDP). PyTorch FSDP, por padrão, fragmentos em todo o conjunto de GPUs em uso. No SMP v2, a biblioteca oferece esse paralelismo de dados fragmentados além do PyTorch FSDP, estendendo a fragmentação PyTorch híbrida (HYBRID_SHARD), que é uma das estratégias de fragmentação fornecidas pelo FSDP:,,,. PyTorch 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 para FSDP. AWS PyTorch

A biblioteca SMP facilita o uso HYBRID_SHARD _HYBRID_SHARD_ZERO2 em qualquer número configurável de GPUs, estendendo o PyTorch FSDP nativo que suporta fragmentação em um único nó () ou em todas as GPUs ()HYBRID_SHARD. FULL_SHARD PyTorch As chamadas FSDP podem permanecer como estão, e você só precisa adicionar o hybrid_shard_degree argumento à configuração SMP, conforme mostrado no exemplo de código a seguir. Você não precisa alterar o valor do sharding_strategy argumento no invólucro do PyTorch FSDP em torno do seu modelo. PyTorch Você pode passar ShardingStrategy.HYBRID_SHARD como valor. Como alternativa, a biblioteca SMP 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 parâmetro. hybrid_shard_degree

Os trechos de código a seguir mostram como adicionar o módulo de inicialização SMP torch.sagemaker.init() ao seu script de treinamento e configurar o dicionário de configuração SMP no formato JSON para o inicializador 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 no PyTorch modelo ou na configuração do PyTorch FSDP. Para obter mais informações sobre o parâmetro hybrid_shard_degree, consulte Parâmetros de configuração do recurso principal do SMP v2.

Dicionário de configuração SMP

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