SMPparâmetros de configuração do recurso principal v2 Referência para o pacote SMP v2 torch.sagemaker Atualize da SMP v1 para a SMP v2

A referência da biblioteca paralela do SageMaker modelo v2

A seguir estão as referências para a biblioteca paralela de SageMaker modelos v2 (SMPv2).

Tópicos

SMPparâmetros de configuração do recurso principal v2
Referência para o pacote SMP v2 torch.sagemaker
Atualize da SMP v1 para a SMP v2

SMPparâmetros de configuração do recurso principal v2

A seguir está uma lista completa de parâmetros para ativar e configurar Principais características da biblioteca de paralelismo de SageMaker modelos v2 o. Eles devem ser escritos em JSON formato e passados para o PyTorch estimador no SageMaker Python SDK ou salvos como um arquivo para. JSON SageMaker HyperPod


{
    "hybrid_shard_degree": Integer,
    "sm_activation_offloading": Boolean,
    "activation_loading_horizon": Integer,
    "fsdp_cache_flush_warnings": Boolean,
    "allow_empty_shards": Boolean,
    "tensor_parallel_degree": Integer,
    "context_parallel_degree": Integer,
    "expert_parallel_degree": Integer,
    "random_seed": Integer
}

hybrid_shard_degree(Inteiro) — Especifica um grau de paralelismo fragmentado. O valor deve ser um número inteiro entre 0 e. world_size O valor padrão é 0.
- Se definido como0, ele volta para a PyTorch implementação nativa e, API no script, quando tensor_parallel_degree é 1. Caso contrário, ele calcula o maior possível hybrid_shard_degree com base em tensor_parallel_degree e. world_size Ao recorrer aos casos de PyTorch FSDP uso nativos, se FULL_SHARD for a estratégia que você usa, ela se fragmenta em todo o cluster deGPUs. Se HYBRID_SHARD ou _HYBRID_SHARD_ZERO2 foi a estratégia, é equivalente hybrid_shard_degree a 8. Quando o paralelismo de tensores está ativado, ele se fragmenta com base na versão revisada. hybrid_shard_degree
- Se definido como1, ele volta para a PyTorch implementação nativa e API para NO_SHARD no script quando tensor_parallel_degree é 1. Caso contrário, é equivalente a qualquer NO_SHARD grupo paralelo de tensores.
- Se definido como um número inteiro entre 2 eworld_size, a fragmentação ocorre no número especificado de. GPUs Se você não configurar sharding_strategy no FSDP script, ele será substituído por. HYBRID_SHARD Se você definir_HYBRID_SHARD_ZERO2, o sharding_strategy que você especificar será usado.
sm_activation_offloading(Boolean) — Especifica se a implementação de descarregamento de SMP ativação deve ser ativada. SeFalse, o descarregamento usa a implementação nativa PyTorch . SeTrue, ele usa a implementação de descarregamento de SMP ativação. Você também precisa usar a PyTorch ativação offload wrapper (torch.distributed.algorithms._checkpoint.checkpoint_wrapper.offload_wrapper) em seu script. Para saber mais, consulte Ativação e descarregamento. O valor padrão é True.
activation_loading_horizon(Integer) — Um número inteiro especificando o tipo de horizonte de descarga de ativação para. FSDP Esse é o número máximo de camadas controladas ou descarregadas cujas entradas podem estar na memória simultaneamente. GPU Para saber mais, consulte Ativação e descarregamento. O valor de entrada deve ser um número inteiro positivo. O valor padrão é 2.
fsdp_cache_flush_warnings(Boolean) — Detecta e avisa se as descargas de cache ocorrem no gerenciador de PyTorch memória, pois elas podem degradar o desempenho computacional. O valor padrão é True.
allow_empty_shards(Boolean) — Se deve permitir fragmentos vazios ao fragmentar tensores se o tensor não for divisível. Essa é uma correção experimental para falhas durante o checkpoint em determinados cenários. Desativar isso remonta ao PyTorch comportamento original. O valor padrão é False.
tensor_parallel_degree(Integer) — Especifica um grau de paralelismo do tensor. O valor deve estar entre 1 world_size e. O valor padrão é 1. Observe que passar um valor maior que 1 não ativa automaticamente o paralelismo de contexto; você também precisa usar o torch.sagemaker.transform API para incluir o modelo em seu script de treinamento. Para saber mais, consulte Paralelismo de tensores.
context_parallel_degree(Integer) — Especifica o grau de paralelismo do contexto. O valor deve estar entre 1 world_size e e deve ser<= hybrid_shard_degree. O valor padrão é 1. Observe que passar um valor maior que 1 não ativa automaticamente o paralelismo de contexto; você também precisa usar o torch.sagemaker.transform API para incluir o modelo em seu script de treinamento. Para saber mais, consulte Paralelismo de contexto.
expert_parallel_degree(Integer) — Especifica um grau de paralelismo especializado. O valor deve estar entre 1 world_size e. O valor padrão é 1. Observe que passar um valor maior que 1 não ativa automaticamente o paralelismo de contexto; você também precisa usar o torch.sagemaker.transform API para incluir o modelo em seu script de treinamento. Para saber mais, consulte Paralelismo especializado.
random_seed(Integer) — Um número inicial para as operações aleatórias em módulos distribuídos por SMP paralelismo tensorial ou paralelismo especializado. Essa semente é adicionada às classificações tensor-parallel ou expert-parallel para definir a semente real para cada classificação. É exclusivo para cada classificação tensor-parallel e expert-parallel. SMPA v2 garante que o número aleatório gerado nas classificações tensor-parallel e expert-parallel corresponda aos casos e, respectivamente. non-tensor-parallelism non-expert-parallelism

Referência para o pacote SMP v2 `torch.sagemaker`

Esta seção é uma referência para o torch.sagemaker pacote fornecido pela SMP v2.

Tópicos

torch.sagemaker.delayed_param.DelayedParamIniter
torch.sagemaker.distributed.checkpoint.state_dict_saver.async_save
torch.sagemaker.distributed.checkpoint.state_dict_saver.maybe_finalize_async_calls
torch.sagemaker.distributed.checkpoint.state_dict_saver.save
torch.sagemaker.distributed.checkpoint.state_dict_loader.load
torch.sagemaker.moe.moe_config.MoEConfig
torch.sagemaker.nn.attn.FlashSelfAttention
torch.sagemaker.nn.attn.FlashGroupedQueryAttention
torch.sagemaker.nn.huggingface.llama_flashattn.LlamaFlashAttention
torch.sagemaker.transform
torch.sagemakerfunções e propriedades do utilitário

`torch.sagemaker.delayed_param.DelayedParamIniter`

E API para se Inicialização atrasada de parâmetros inscrever em um PyTorch modelo.


class torch.sagemaker.delayed_param.DelayedParamIniter(
    model: nn.Module,
    init_method_using_config : Callable = None,
    verbose: bool = False,
)

Parâmetros

model(nn.Module) — Um PyTorch modelo para empacotar e aplicar a funcionalidade de inicialização retardada de parâmetros da SMP v2.
init_method_using_config(Callable) — Se você usar a implementação paralela de tensor da SMP v2 ou suportadaModelos Hugging Face Transformer compatíveis com o paralelismo tensorial SMP, mantenha esse parâmetro no valor padrão, que é. None Por padrão, DelayedParamIniter API ele descobre como inicializar o modelo fornecido corretamente. Para qualquer outro modelo, você precisa criar uma função de inicialização de parâmetros personalizada e adicioná-la ao seu script. O trecho de código a seguir é a init_method_using_config função padrão que a SMP v2 implementou para o. Modelos Hugging Face Transformer compatíveis com o paralelismo tensorial SMP Use o seguinte trecho de código como referência para criar sua própria função de configuração de inicialização, adicioná-la ao seu script e passá-la para o init_method_using_config parâmetro do. SMP DelayedParamIniter API
```
from torch.sagemaker.utils.module_utils import empty_module_params, move_buffers_to_device

# Define a custom init config function.
def custom_init_method_using_config(module):
    d = torch.cuda.current_device()
    empty_module_params(module, device=d)
    if isinstance(module, (nn.Linear, Conv1D)):
        module.weight.data.normal_(mean=0.0, std=config.initializer_range)
        if module.bias is not None:
            module.bias.data.zero_()
    elif isinstance(module, nn.Embedding):
        module.weight.data.normal_(mean=0.0, std=config.initializer_range)
        if module.padding_idx is not None:
            module.weight.data[module.padding_idx].zero_()
    elif isinstance(module, nn.LayerNorm):
        module.weight.data.fill_(1.0)
        module.bias.data.zero_()
    elif isinstance(module, LlamaRMSNorm):
        module.weight.data.fill_(1.0)
    move_buffers_to_device(module, device=d)

delayed_initer = DelayedParamIniter(model, init_method_using_config=custom_init_method_using_config)
```
Para obter mais informações sobre as torch.sagemaker.module_util funções no trecho de código anterior, consulte. torch.sagemakerfunções e propriedades do utilitário
verbose(Boolean) — Se é necessário ativar um registro mais detalhado durante a inicialização e a validação. O valor padrão é False.

Métodos

get_param_init_fn()— Retorna a função de inicialização do parâmetro que você pode passar para o param_init_fn argumento da classe PyTorch FSDP wrapper.
get_post_param_init_fn()— Retorna a função de inicialização do parâmetro que você pode passar para o post_param_init_fn argumento da classe PyTorch FSDP wrapper. Isso é necessário quando você amarra pesos no modelo. O modelo deve implementar o métodotie_weights. Para obter mais informações, consulte as notas sobre peso vinculadoInicialização atrasada de parâmetros.
count_num_params(module: nn.Module, *args: Tuple[nn.Parameter]) — Rastreia quantos parâmetros estão sendo inicializados pela função de inicialização de parâmetros. Isso ajuda a implementar o validate_params_and_buffers_inited método a seguir. Normalmente, você não precisa chamar essa função explicitamente, porque o validate_params_and_buffers_inited método chama implicitamente esse método no back-end.
validate_params_and_buffers_inited(enabled: bool=True) — Este é um gerenciador de contexto que ajuda a validar se o número de parâmetros inicializados corresponde ao número total de parâmetros no modelo. Ele também valida que todos os parâmetros e buffers estão agora em GPU dispositivos em vez de meta-dispositivos. Isso aumenta AssertionErrors se essas condições não forem atendidas. Esse gerenciador de contexto é apenas opcional e você não precisa usá-lo para inicializar parâmetros.

`torch.sagemaker.distributed.checkpoint.state_dict_saver.async_save`

Entrada API para salvamento assíncrono. Use esse método para salvar um de forma state_dict assíncrona em um especificado. checkpoint_id


def async_save(
    state_dict: STATE_DICT_TYPE,
    *,
    checkpoint_id: Union[str, os.PathLike, None] = None,
    storage_writer: Optional[StorageWriter] = None,
    planner: Optional[SavePlanner] = None,
    process_group: Optional[dist.ProcessGroup] = None,
    coordinator_rank: int = 0,
    queue : AsyncCallsQueue = None,
    sharded_strategy: Union[SaveShardedStrategy, Tuple[str, int], None] = None,
    wait_error_handling: bool = True,
    force_check_all_plans: bool = True,
    s3_region: Optional[str] = None,
    s3client_config: Optional[S3ClientConfig] = None
) -> None:

Parâmetros

state_dict(dict) - Obrigatório. O ditado estadual para salvar.
checkpoint_id(str) - Obrigatório. O caminho de armazenamento no qual salvar os pontos de verificação.
storage_writer(StorageWriter) - Opcional. Uma instância de StorageWriterin PyTorch para realizar operações de gravação. Se isso não for especificado, a configuração padrão de será StorageWriterusada.
planner(SavePlanner) - Opcional. Uma instância de SavePlannerentrada PyTorch. Se isso não for especificado, a configuração padrão de será SavePlannerusada.
process_group(ProcessGroup) - Opcional. O grupo de processos no qual trabalhar. SeNone, o grupo de processos padrão (global) for usado.
coordinator_rank(int) - Opcional. A classificação do coordenador ao realizar operadores de comunicação coletiva, comoAllReduce.
queue(AsyncRequestQueue) - Opcional. O agendador assíncrono a ser usado. Por padrão, ele usa o parâmetro globalDEFAULT_ASYNC_REQUEST_QUEUE.
sharded_strategy(PyTorchDistSaveShardedStrategy) - Opcional. A estratégia fragmentada a ser usada para salvar pontos de verificação. Se isso não for especificado, torch.sagemaker.distributed.checkpoint.state_dict_saver.PyTorchDistSaveShardedStrategy é usado por padrão.
wait_error_handling(bool) - Opcional. Um sinalizador que especifica se deve esperar que todas as classificações concluam o tratamento de erros. O valor padrão é True.
force_check_all_plans(bool) - Opcional. Um sinalizador que determina se os planos devem ser sincronizados à força em todos os níveis, mesmo no caso de uma falha no cache. O valor padrão é True.
s3_region(str) - Opcional. A região em que o bucket do S3 está localizado. Se não for especificada, a região será inferida docheckpoint_id.
s3client_config(S3ClientConfig) - Opcional. A classe de dados que expõe parâmetros configuráveis para o cliente S3. Se não for fornecida, a configuração padrão do S3 ClientConfig será usada. O part_size parâmetro é definido como 64 MB por padrão.

`torch.sagemaker.distributed.checkpoint.state_dict_saver.maybe_finalize_async_calls`

Essa função permite que um processo de treinamento monitore várias solicitações assíncronas a serem feitas.


def maybe_finalize_async_calls(
    blocking=True, 
    process_group=None
) -> List[int]:

Parâmetros

blocking(bool) - Opcional. SeTrue, ele aguardará até que todas as solicitações ativas sejam concluídas. Caso contrário, ele finaliza somente as solicitações assíncronas que já foram concluídas. O valor padrão é True.
process_group(ProcessGroup) - Opcional. O grupo de processos no qual operar. Se definido comoNone, o grupo de processos padrão (global) será utilizado.

Devoluções

Uma lista contendo os índices de chamadas assíncronas foi finalizada com sucesso.

`torch.sagemaker.distributed.checkpoint.state_dict_saver.save`

Use esse método para salvar um de state_dict forma síncrona em um especificadocheckpoint_id.


def save(
    state_dict: STATE_DICT_TYPE,
    *,
    checkpoint_id: Union[str, os.PathLike, None] = None,
    storage_writer: Optional[StorageWriter] = None,
    planner: Optional[SavePlanner] = None,
    process_group: Optional[dist.ProcessGroup] = None,
    coordinator_rank: int = 0,
    wait_error_handling: bool = True,
    force_check_all_plans: bool = True,
    s3_region: Optional[str] = None,
    s3client_config: Optional[S3ClientConfig] = None
) -> None:

Parâmetros

state_dict(dict) - Obrigatório. O ditado estadual para salvar.
checkpoint_id(str) - Obrigatório. O caminho de armazenamento no qual salvar os pontos de verificação.
storage_writer(StorageWriter) - Opcional. Uma instância de StorageWriterin PyTorch para realizar operações de gravação. Se isso não for especificado, a configuração padrão de será StorageWriterusada.
planner(SavePlanner) - Opcional. Uma instância de SavePlannerentrada PyTorch. Se isso não for especificado, a configuração padrão de será SavePlannerusada.
process_group(ProcessGroup) - Opcional. O grupo de processos no qual trabalhar. SeNone, o grupo de processos padrão (global) for usado.
coordinator_rank(int) - Opcional. A classificação do coordenador ao realizar operadores de comunicação coletiva, comoAllReduce.
wait_error_handling(bool) - Opcional. Um sinalizador que especifica se deve esperar que todas as classificações concluam o tratamento de erros. O valor padrão é True.
force_check_all_plans(bool) - Opcional. Um sinalizador que determina se os planos devem ser sincronizados à força em todos os níveis, mesmo no caso de uma falha no cache. O valor padrão é True.
s3_region(str) - Opcional. A região em que o bucket do S3 está localizado. Se não for especificada, a região será inferida docheckpoint_id.
s3client_config(S3ClientConfig) - Opcional. A classe de dados que expõe parâmetros configuráveis para o cliente S3. Se não for fornecida, a configuração padrão do S3 ClientConfig será usada. O part_size parâmetro é definido como 64 MB por padrão.

`torch.sagemaker.distributed.checkpoint.state_dict_loader.load`

Carregue o dicionário de estados de um modelo distribuído (state_dict).


def load(
    state_dict: Dict[str, Any],
    *,
    checkpoint_id: Union[str, os.PathLike, None] = None,
    storage_reader: Optional[StorageReader] = None,
    planner: Optional[LoadPlanner] = None,
    process_group: Optional[dist.ProcessGroup] = None,
    check_keys_matched: bool = True,
    coordinator_rank: int = 0,
    s3_region: Optional[str] = None,
    s3client_config: Optional[S3ClientConfig] = None
) -> None:

Parâmetros

state_dict(dict) - Obrigatório. O state_dict para carregar.
checkpoint_id(str) - Obrigatório. O ID de um posto de controle. O significado do checkpoint_id depende do armazenamento. Pode ser um caminho para uma pasta ou para um arquivo. Também pode ser uma chave se o armazenamento for um armazenamento de valores-chave.
storage_reader(StorageReader) - Opcional. Uma instância de StorageReaderin PyTorch para realizar operações de leitura. Se não for especificado, o ponto de verificação distribuído inferirá automaticamente o leitor com base no. checkpoint_id Se também checkpoint_id estiverNone, um erro de exceção será gerado.
planner(StorageReader) - Opcional. Uma instância de LoadPlannerentrada PyTorch. Se não for especificada, a configuração padrão de LoadPlannerserá usada.
check_keys_matched(bool) - Opcional. Se ativado, verifica se as state_dict chaves de todas as classificações são combinadas usandoAllGather.
s3_region(str) - Opcional. A região em que o bucket do S3 está localizado. Se não for especificada, a região será inferida docheckpoint_id.
s3client_config(S3ClientConfig) - Opcional. A classe de dados que expõe parâmetros configuráveis para o cliente S3. Se não for fornecida, a configuração padrão do S3 ClientConfig será usada. O part_size parâmetro é definido como 64 MB por padrão.

`torch.sagemaker.moe.moe_config.MoEConfig`

Uma classe de configuração para configurar a SMP implementação do Mixture-of-Experts (MoE). Você pode especificar valores de configuração do MoE por meio dessa classe e passá-los para a torch.sagemaker.transformAPIchamada. Para saber mais sobre o uso dessa classe para treinar modelos MoE, consulteParalelismo especializado.


class torch.sagemaker.moe.moe_config.MoEConfig(
    smp_moe=True,
    random_seed=12345,
    moe_load_balancing="sinkhorn",
    global_token_shuffle=False,
    moe_all_to_all_dispatcher=True,
    moe_aux_loss_coeff=0.001,
    moe_z_loss_coeff=0.001
)

Parâmetros

smp_moe(Boolean) - Se deve usar a SMP implementação -do MoE. O valor padrão é True.
random_seed(Integer) - Um número inicial para as operações aleatórias em módulos distribuídos paralelos especializados. Essa semente é adicionada à classificação paralela de especialistas para definir a semente real para cada classificação. É exclusivo para cada classificação paralela de especialistas. O valor padrão é 12345.
moe_load_balancing(String) - Especifique o tipo de balanceamento de carga do roteador MoE. As opções válidas são aux_loss sinkhornbalanced,, none e. O valor padrão é sinkhorn.
global_token_shuffle(Boolean) - Se os tokens devem ser misturados entre as classificações do EP dentro do mesmo grupo de EP. O valor padrão é False.
moe_all_to_all_dispatcher(Boolean) - Se deve usar o all-to-all dispatcher para as comunicações no MoE. O valor padrão é True.
moe_aux_loss_coeff(Flutuar) - Um coeficiente para perda de balanceamento de carga auxiliar. O valor padrão é 0.001.
moe_z_loss_coeff(Float) - Coeficiente de perda z. O valor padrão é 0.001.

`torch.sagemaker.nn.attn.FlashSelfAttention`

E API para usar FlashAttention com a SMP v2.


class torch.sagemaker.nn.attn.FlashSelfAttention(
   attention_dropout_prob: float = 0.0,
   scale: Optional[float] = None,
   triton_flash_attention: bool = False,
   use_alibi: bool = False,
)

Parâmetros

attention_dropout_prob(float) — A probabilidade de abandono escolar a ser aplicada à atenção. O valor padrão é 0.0.
scale(float) — Se aprovado, esse fator de escala é aplicado para softmax. Se definido como None (que também é o valor padrão), o fator de escala é1 / sqrt(attention_head_size). O valor padrão é None.
triton_flash_attention(bool) — Se aprovada, a implementação de atenção instantânea do Triton é usada. Isso é necessário para apoiar a atenção com vieses lineares (ALiBi) (consulte o use_alibi parâmetro a seguir). Essa versão do kernel não suporta o dropout. O valor padrão é False.

use_alibi(bool) — Se aprovado, ele ativa a Atenção com vieses lineares (ALiBi) usando a máscara fornecida. Ao usarALiBi, ele precisa de uma máscara de atenção preparada da seguinte forma. O valor padrão é False.


def generate_alibi_attn_mask(attention_mask, batch_size, seq_length, 
    num_attention_heads, alibi_bias_max=8):
    device, dtype = attention_mask.device, attention_mask.dtype
    alibi_attention_mask = torch.zeros(
        1, num_attention_heads, 1, seq_length, dtype=dtype, device=device
    )

    alibi_bias = torch.arange(1 - seq_length, 1, dtype=dtype, device=device).view(
        1, 1, 1, seq_length
    )
    m = torch.arange(1, num_attention_heads + 1, dtype=dtype, device=device)
    m.mul_(alibi_bias_max / num_attention_heads)
    alibi_bias = alibi_bias * (1.0 / (2 ** m.view(1, num_attention_heads, 1, 1)))

    alibi_attention_mask.add_(alibi_bias)
    alibi_attention_mask = alibi_attention_mask[..., :seq_length, :seq_length]
    if attention_mask is not None and attention_mask.bool().any():
        alibi_attention_mask.masked_fill(
            attention_mask.bool().view(batch_size, 1, 1, seq_length), float("-inf")
        )

    return alibi_attention_mask

Métodos

forward(self, qkv, attn_mask=None, causal=False, cast_dtype=None, layout="b h s d")— Uma função regular PyTorch do módulo. Quando a module(x) é chamado, SMP executa essa função automaticamente.
- qkv— torch.Tensor da seguinte forma: (batch_size x seqlen x (3 x num_heads) x head_size) ou(batch_size, (3 x num_heads) x seqlen x head_size), uma tupla de torch.Tensors cada uma das quais pode ter a forma(batch_size x seqlen x num_heads x head_size), ou(batch_size x num_heads x seqlen x head_size). Um argumento de layout apropriado deve ser passado com base na forma.
- attn_mask— torch.Tensor do seguinte formulário(batch_size x 1 x 1 x seqlen). Para ativar esse parâmetro de máscara de atenção, é necessário triton_flash_attention=True use_alibi=True e. Para saber como gerar uma máscara de atenção usando esse método, consulte os exemplos de código emFlashAttention. O valor padrão é None.
- causal— Quando definido comoFalse, que é o valor padrão do argumento, nenhuma máscara é aplicada. Quando definido comoTrue, o forward método usa a máscara triangular inferior padrão. O valor padrão é False.
- cast_dtype— Quando configurado para um determinadodtype, ele converte os qkv tensores para aquele dtype anteriorattn. Isso é útil para implementações como o modelo Hugging Face GPT Transformer -NEOx, que tem e com incorporações rotativas posteriores. q k fp32 Se definido comoNone, nenhum molde será aplicado. O valor padrão é None.
- layout(string) — Os valores disponíveis são b h s d oub s h d. Isso deve ser definido para o layout dos qkv tensores passados, para que as transformações apropriadas possam ser aplicadas. attn O valor padrão é b h s d.

Devoluções

Um single torch.Tensor com forma(batch_size x num_heads x seq_len x head_size).

`torch.sagemaker.nn.attn.FlashGroupedQueryAttention`

E API para usar FlashGroupedQueryAttention com a SMP v2. Para saber mais sobre o uso dissoAPI, consulteUse FlashAttention kernels para atenção de consultas agrupadas.


class torch.sagemaker.nn.attn.FlashGroupedQueryAttention(
    attention_dropout_prob: float = 0.0,
    scale: Optional[float] = None,
)

Parâmetros

attention_dropout_prob(float) — A probabilidade de abandono escolar a ser aplicada à atenção. O valor padrão é 0.0.
scale(float) — Se aprovado, esse fator de escala é aplicado para softmax. Se definido comoNone, 1 / sqrt(attention_head_size) é usado como fator de escala. O valor padrão é None.

Métodos

forward(self, q, kv, causal=False, cast_dtype=None, layout="b s h d")— Uma função regular PyTorch do módulo. Quando a module(x) é chamado, SMP executa essa função automaticamente.
- q— torch.Tensor da seguinte forma (batch_size x seqlen x num_heads x head_size) ou(batch_size x num_heads x seqlen x head_size). O argumento de layout apropriado deve ser passado com base na forma.
- kv— torch.Tensor da seguinte forma (batch_size x seqlen x (2 x num_heads) x head_size) ou(batch_size, (2 x num_heads) x seqlen x head_size), ou uma tupla de dois torch.Tensor s, cada um dos quais pode ter a forma (batch_size x seqlen x num_heads x head_size) ou(batch_size x num_heads x seqlen x head_size). layoutO argumento apropriado também deve ser passado com base na forma.
- causal— Quando definido comoFalse, que é o valor padrão do argumento, nenhuma máscara é aplicada. Quando definido comoTrue, o forward método usa a máscara triangular inferior padrão. O valor padrão é False.
- cast_dtype— Quando definido para um determinado dtype, ele converte os qkv tensores para esse dtype antes. attn Isso é útil para implementações como Hugging Face GPT Transformers -NeoX, que tem incorporações rotativas posteriores. q,k fp32 Se definido comoNone, nenhum molde será aplicado. O valor padrão é None.
- layout (string) — Os valores disponíveis são "b h s d" ou"b s h d". Isso deve ser definido para o layout dos qkv tensores passados, para que as transformações apropriadas possam ser aplicadas. attn O valor padrão é "b h s d".

Devoluções

Retorna um single torch.Tensor (batch_size x num_heads x seq_len x head_size) que representa a saída do cálculo da atenção.

`torch.sagemaker.nn.huggingface.llama_flashattn.LlamaFlashAttention`

E API isso suporta FlashAttention o modelo Llama. Isso API usa o torch.sagemaker.nn.attn.FlashGroupedQueryAttention API em baixo nível. Para saber como usar isso, consulteUse FlashAttention kernels para atenção de consultas agrupadas.


class torch.sagemaker.nn.huggingface.llama_flashattn.LlamaFlashAttention(
    config: LlamaConfig
)

Parâmetros

config— Uma FlashAttention configuração para o modelo Llama.

Métodos

forward(self, hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache)
- hidden_states(torch.Tensor) — Estados ocultos de um tensor na forma de(batch_size x seq_len x num_heads x head_size).
- attention_mask(torch.LongTensor) — Máscara para evitar prestar atenção ao preenchimento de índices de tokens na forma de. (batch_size x seqlen) O valor padrão é None.
- position_ids(torch.LongTensor) — Quando não estáNone, está na forma de(batch_size x seqlen), indicando os índices das posições de cada token de sequência de entrada nas incorporações da posição. O valor padrão é None.
- past_key_value(Cache) — Estados ocultos pré-computados (chave e valores nos blocos de autoatenção e nos blocos de atenção cruzada). O valor padrão é None.
- output_attentions(bool) — Indica se os tensores de atenção de todas as camadas de atenção devem ser retornados. O valor padrão é False.
- use_cache(bool) — Indica se os estados do valor da past_key_values chave devem ser retornados. O valor padrão é False.

Devoluções

Retorna um single torch.Tensor (batch_size x num_heads x seq_len x head_size) que representa a saída do cálculo da atenção.

`torch.sagemaker.transform`

SMPA v2 fornece isso torch.sagemaker.transform() API para transformar os modelos do Hugging Face Transformer em implementações de SMP modelos e habilitar o paralelismo do tensor. SMP


torch.sagemaker.transform(
    model: nn.Module, 
    device: Optional[torch.device] = None, 
    dtype: Optional[torch.dtype] = None, 
    config: Optional[Dict] = None, 
    load_state_dict_from_rank0: bool = False
)

SMPA v2 mantém as políticas de transformação para o Modelos Hugging Face Transformer compatíveis com o paralelismo tensorial SMP convertendo a configuração dos modelos do Hugging Face Transformer na configuração do transformador. SMP

Parâmetros

model(torch.nn.Module) — Um modelo Modelos Hugging Face Transformer compatíveis com o paralelismo tensorial SMP para transformar e aplicar o recurso de paralelismo tensorial da biblioteca. SMP
device(torch.device) — Se aprovado, um novo modelo é criado neste dispositivo. Se o módulo original tiver algum parâmetro no meta-dispositivo (consulteInicialização atrasada de parâmetros), o módulo transformado também será criado no meta-dispositivo, ignorando o argumento passado aqui. O valor padrão é None.
dtype(torch.dtype) — Se aprovado, define isso como o gerenciador de contexto dtype para a criação do modelo e cria um modelo com esse dtype. Normalmente, isso é desnecessário, pois queremos criar o modelo fp32 ao usar MixedPrecision e fp32 é o dtype in PyTorch padrão. O valor padrão é None.
config(dict) — Este é um dicionário para configurar o SMP transformador. O valor padrão é None.
load_state_dict_from_rank0(Boolean) — Por padrão, esse módulo cria uma nova instância do modelo com novos pesos. Quando esse argumento é definido comoTrue, SMP tenta carregar o dicionário de estados do PyTorch modelo original da 0ª classificação em um modelo transformado para o grupo paralelo de tensores do qual a 0ª classificação faz parte. Quando definido comoTrue, a classificação 0 não pode ter nenhum parâmetro no meta-dispositivo. Somente o primeiro grupo paralelo de tensores preenche os pesos da 0ª classificação após essa chamada de transformação. Você precisa definir como sync_module_states True no FSDP wrapper para obter esses pesos do primeiro grupo paralelo de tensores para todos os outros processos. Com isso ativado, a SMP biblioteca carrega o dicionário de estados do modelo original. A SMP biblioteca pega o state_dict do modelo antes da transformação, o converte para corresponder à estrutura do modelo transformado, o fragmenta para cada classificação paralela do tensor, comunica esse estado da 0ª classificação para outras classificações no grupo paralelo de tensores do qual a 0ª classificação faz parte e o carrega. O valor padrão é False.

Devoluções

Retorna um modelo transformado com o qual você pode embrulhar PyTorch FSDP. Quando load_state_dict_from_rank0 definido comoTrue, o grupo paralelo de tensores que envolve a classificação 0 tem pesos carregados do dicionário de estado original na classificação 0. Ao usar Inicialização atrasada de parâmetros no modelo original, somente essas classificações têm os tensores reais ativados CPUs para os parâmetros e buffers do modelo transformado. O resto das classificações continuam com os parâmetros e buffers no meta-dispositivo para economizar memória.

`torch.sagemaker`funções e propriedades do utilitário

Funções do utilitário torch.sagemaker

torch.sagemaker.init(config: Optional[Union[str, Dict[str, Any]]] = None) -> None— Inicializa o trabalho PyTorch de treinamento comSMP.
torch.sagemaker.is_initialized() -> bool— Verifica se o trabalho de treinamento foi inicializado comSMP. Ao voltar para o nativo PyTorch enquanto o trabalho é inicializado comSMP, algumas das propriedades não são relevantes e se tornamNone, conforme indicado na lista de propriedades a seguir.
torch.sagemaker.utils.module_utils.empty_module_params(module: nn.Module, device: Optional[torch.device] = None, recurse: bool = False) -> nn.Module— Cria parâmetros vazios no dado, device se houver, e pode ser recursivo para todos os módulos aninhados, se especificado.
torch.sagemaker.utils.module_utils.move_buffers_to_device(module: nn.Module, device: torch.device, recurse: bool = False) -> nn.Module— Move os buffers do módulo para o determinado device e pode ser recursivo para todos os módulos aninhados, se especificado.

Propriedades

torch.sagemaker.statecontém várias propriedades úteis após a inicialização de SMP withtorch.sagemaker.init.

torch.sagemaker.state.hybrid_shard_degree(int) — O grau de paralelismo de dados fragmentados, uma cópia da entrada do usuário na SMP configuração passada para. torch.sagemaker.init() Para saber mais, consulte Comece com a biblioteca de paralelismo de SageMaker modelos v2.
torch.sagemaker.state.rank(int) — A classificação global do dispositivo, na faixa de[0, world_size).
torch.sagemaker.state.rep_rank_process_group(torch.distributed.ProcessGroup) — O grupo de processos que inclui todos os dispositivos com a mesma classificação de replicação. Observe a diferença sutil, mas fundamental, comtorch.sagemaker.state.tp_process_group. Ao voltar para o nativo PyTorch, ele retornaNone.
torch.sagemaker.state.tensor_parallel_degree(int) — O grau de paralelismo do tensor, uma cópia da entrada do usuário na SMP configuração passada para. torch.sagemaker.init() Para saber mais, consulte Comece com a biblioteca de paralelismo de SageMaker modelos v2.
torch.sagemaker.state.tp_size(int) — Um alias paratorch.sagemaker.state.tensor_parallel_degree.
torch.sagemaker.state.tp_rank(int) — A classificação do paralelismo do tensor para o dispositivo na faixa de[0, tp_size), determinada pelo grau de paralelismo do tensor e pelo mecanismo de classificação.
torch.sagemaker.state.tp_process_group(torch.distributed.ProcessGroup) — O grupo de processos paralelos de tensores, incluindo todos os dispositivos com a mesma classificação em outras dimensões (por exemplo, paralelismo e replicação de dados fragmentados), mas classificações paralelas de tensores exclusivos. Ao voltar para o nativo PyTorch, ele retornaNone.
torch.sagemaker.state.world_size(int) — O número total de dispositivos usados no treinamento.

Atualize da SMP v1 para a SMP v2

Para passar da SMP v1 para a SMP v2, você deve fazer alterações no script para remover a SMP v1 APIs e aplicar a v2. SMP APIs Em vez de começar com seu script SMP v1, recomendamos que você comece com um PyTorch FSDP script e siga as instruções emComece com a biblioteca de paralelismo de SageMaker modelos v2.

Para trazer modelos SMP v1 para SMP v2, na SMP v1 você deve coletar o dicionário de estado do modelo completo e aplicar as funções de tradução no dicionário de estado do modelo para convertê-lo no formato de ponto de verificação do modelo Hugging Face Transformers. Então, na SMP v2, conforme discutido emPonto de verificação usando SMP, você pode carregar os pontos de verificação do modelo Hugging Face Transformers e continuar usando o ponto de verificação com a v2. PyTorch APIs SMP Para usar SMP com seu PyTorch FSDP modelo, certifique-se de migrar para a SMP v2 e fazer alterações no script de treinamento para uso PyTorch FSDP e em outros recursos mais recentes.


import smdistributed.modelparallel.torch as smp

# Create model
model = ...
model = smp.DistributedModel(model)

# Run training
...

# Save v1 full checkpoint
if smp.rdp_rank() == 0:
    model_dict = model.state_dict(gather_to_rank0=True) # save the full model
    # Get the corresponding translation function in smp v1 and translate
    if model_type == "gpt_neox":
        from smdistributed.modelparallel.torch.nn.huggingface.gptneox import translate_state_dict_to_hf_gptneox
        translated_state_dict = translate_state_dict_to_hf_gptneox(state_dict, max_seq_len=None)
    
    # Save the checkpoint
    checkpoint_path = "checkpoint.pt"
    if smp.rank() == 0:
        smp.save(
            {"model_state_dict": translated_state_dict},
            checkpoint_path,
            partial=False,
        )

Para encontrar as funções de tradução disponíveis na SMP v1, consulteSuporte para modelos Hugging Face Transformer.

Para obter instruções sobre como salvar e carregar pontos de verificação de modelos na SMP v2, consulte. Ponto de verificação usando SMP

Atenção O Javascript está desativado ou não está disponível no seu navegador.

Para usar a documentação da AWS, o Javascript deve estar ativado. Consulte as páginas de Ajuda do navegador para obter instruções.

Convenções do documento

Práticas recomendadas

Notas de release

A referência da biblioteca paralela do SageMaker modelo v2

Tópicos

SMPparâmetros de configuração do recurso principal v2

Referência para o pacote SMP v2 torch.sagemaker

Tópicos

torch.sagemaker.delayed_param.DelayedParamIniter

torch.sagemaker.distributed.checkpoint.state_dict_saver.async_save

torch.sagemaker.distributed.checkpoint.state_dict_saver.maybe_finalize_async_calls

torch.sagemaker.distributed.checkpoint.state_dict_saver.save

torch.sagemaker.distributed.checkpoint.state_dict_loader.load

torch.sagemaker.moe.moe_config.MoEConfig

torch.sagemaker.nn.attn.FlashSelfAttention

torch.sagemaker.nn.attn.FlashGroupedQueryAttention

torch.sagemaker.nn.huggingface.llama_flashattn.LlamaFlashAttention

torch.sagemaker.transform

torch.sagemakerfunções e propriedades do utilitário

Funções do utilitário torch.sagemaker

Propriedades

Atualize da SMP v1 para a SMP v2

Referência para o pacote SMP v2 `torch.sagemaker`

`torch.sagemaker.delayed_param.DelayedParamIniter`

`torch.sagemaker.distributed.checkpoint.state_dict_saver.async_save`

`torch.sagemaker.distributed.checkpoint.state_dict_saver.maybe_finalize_async_calls`

`torch.sagemaker.distributed.checkpoint.state_dict_saver.save`

`torch.sagemaker.distributed.checkpoint.state_dict_loader.load`

`torch.sagemaker.moe.moe_config.MoEConfig`

`torch.sagemaker.nn.attn.FlashSelfAttention`

`torch.sagemaker.nn.attn.FlashGroupedQueryAttention`

`torch.sagemaker.nn.huggingface.llama_flashattn.LlamaFlashAttention`

`torch.sagemaker.transform`

`torch.sagemaker`funções e propriedades do utilitário