Referenz für die SageMaker-Modellparallelitätsbibliothek v2 - Amazon SageMaker KI

Die vorliegende Übersetzung wurde maschinell erstellt. Im Falle eines Konflikts oder eines Widerspruchs zwischen dieser übersetzten Fassung und der englischen Fassung (einschließlich infolge von Verzögerungen bei der Übersetzung) ist die englische Fassung maßgeblich.

Referenz für die SageMaker-Modellparallelitätsbibliothek v2

Im Folgenden finden Sie Referenzen für die SageMaker-Modellparallelitätsbibliothek v2 (SMP v2).

Konfigurationsparameter für die Kernfunktionen von SMP v2

Nachfolgend finden Sie eine vollständige Liste der Parameter zur Aktivierung und Konfiguration von Kernfunktionen der SageMaker Modellparallelitätsbibliothek v2. Diese müssen im JSON-Format geschrieben und an den PyTorch-Schätzer im SageMaker Python SDK übergeben oder als JSON-Datei für SageMaker HyperPod gespeichert werden.

{ "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 (Ganzzahl) – Gibt einen Parallelitätsgrad fragmentierter Daten an. Der Wert muss eine Ganzzahl zwischen 0 und world_size sein. Der Standardwert ist 0.

    • Wenn der Wert auf 0 gesetzt ist, wird auf die native PyTorch-Implementierung und API im Skript zurückgegriffen, wenn der Wert für tensor_parallel_degree 1 ist. Andernfalls wird der größtmögliche Wert für tensor_parallel_degree auf der Grundlage von hybrid_shard_degree und world_size berechnet. Wenn Sie auf die nativen PyTorch-FSDP-Anwendungsfälle zurückgreifen und FULL_SHARD als Strategie verwenden, werden die Daten auf dem gesamten Cluster an GPUs fragmentiert. Wenn als Strategie HYBRID_SHARD oder _HYBRID_SHARD_ZERO2 verwendet wird, entspricht der hybrid_shard_degree 8. Wenn die Tensorparallelität aktiviert ist, basiert die Fragmentierung auf dem überarbeiteten hybrid_shard_degree.

    • Wenn der Wert auf 1 gesetzt ist, wird auf die native PyTorch-Implementierung und API für NO_SHARD im Skript zurückgegriffen, wenn der Wert für tensor_parallel_degree 1 ist. Andernfalls entspricht der Wert NO_SHARD innerhalb einer beliebigen Tensorparallelgruppe.

    • Wenn der Wert auf eine Ganzzahl zwischen 2 und world_size gesetzt ist, erfolgt das Sharding für die angegebene Anzahl von GPUs. Wenn Sie nicht sharding_strategy im FSDP-Skript einrichten, wird es in HYBRID_SHARD überschrieben. Wenn Sie _HYBRID_SHARD_ZERO2 festlegen, wird die von Ihnen angegebene sharding_strategy verwendet.

  • sm_activation_offloading (Boolean) – Legt fest, ob die Implementierung der SMP-Aktivierungsauslagerung aktiviert werden soll. Wenn False, wird bei der Auslagerung die native PyTorch-Implementierung verwendet. Wenn True, wird die Implementierung der SMP-Aktivierungsauslagerung verwendet. Sie müssen auch den PyTorch-Aktivierungsauslagerungs-Wrapper (torch.distributed.algorithms._checkpoint.checkpoint_wrapper.offload_wrapper) in Ihrem Skript verwenden. Weitere Informationen hierzu finden Sie unter Aktivierung, Entladung. Der Standardwert ist True.

  • activation_loading_horizon (Ganzzahl) – Eine Ganzzahl, die die Art von Aktivierungsauslagerungshorizont für FSDP angibt. Dies ist die maximale Anzahl von Prüfpunkt- oder Auslagerungsebenen, deren Eingaben sich gleichzeitig im GPU-Speicher befinden können. Weitere Informationen hierzu finden Sie unter Aktivierung, Entladung. Der Eingabewert muss eine positive Ganzzahl sein. Der Standardwert ist 2.

  • fsdp_cache_flush_warnings (Boolean) – Erkennt und warnt, wenn Cache-Leerungen im PyTorch-Speichermanager auftreten, da sie die Rechenleistung beeinträchtigen können. Der Standardwert ist True.

  • allow_empty_shards (Boolean) – Ob beim Tensor-Sharding leere Shards zulässig sind, wenn der Tensor nicht teilbar ist. Dies ist eine experimentelle Lösung für Abstürze beim Checkpointing in bestimmten Szenarien. Das Deaktivieren dieser Option führt zum ursprünglichen PyTorch-Verhalten. Der Standardwert ist False.

  • tensor_parallel_degree (Ganzzahl) – Gibt einen Tensor-Parallelitätsgrad an. Dieser Wert muss zwischen 1 und world_size liegen. Der Standardwert ist 1. Beachten Sie, dass die Übergabe eines Werts größer als 1 die Kontextparallelität nicht automatisch aktiviert. Sie müssen auch die torch.sagemaker.transform-API verwenden, um das Modell in Ihr Trainingsskript einzubinden. Weitere Informationen hierzu finden Sie unter Tensor-Parallelität.

  • context_parallel_degree (Ganzzahl) – Gibt den Grad der Kontextparallelität an. Der Wert muss zwischen 1 und world_size liegen und muss der <= hybrid_shard_degree sein. Der Standardwert ist 1. Beachten Sie, dass die Übergabe eines Werts größer als 1 die Kontextparallelität nicht automatisch aktiviert. Sie müssen auch die torch.sagemaker.transform-API verwenden, um das Modell in Ihr Trainingsskript einzubinden. Weitere Informationen hierzu finden Sie unter Kontextparallelität.

  • expert_parallel_degree (Ganzzahl) – Gibt einen Grad der Expertenparallelität an. Dieser Wert muss zwischen 1 und world_size liegen. Der Standardwert ist 1. Beachten Sie, dass die Übergabe eines Werts größer als 1 die Kontextparallelität nicht automatisch aktiviert. Sie müssen auch die torch.sagemaker.transform-API verwenden, um das Modell in Ihr Trainingsskript einzubinden. Weitere Informationen hierzu finden Sie unter Expertenparallelität.

  • random_seed (Ganzzahl) – Eine Seed-Zahl für zufällige Operationen in verteilten Modulen nach SMP-Tensorparallelität oder Expertenparallelität. Dieser Seed-Wert wird zu tensorparallelen oder expertenparallelen Rängen hinzugefügt, um den tatsächlichen Seed-Wert für jeden Rang festzulegen. Er ist für jeden tensorparallelen und expertenparallelen Rang einzigartig. SMP v2 stellt sicher, dass die Zufallszahl, die für tensorparallele und expertenparallele Ränge generiert wird, mit den Fällen der Nicht-Tensorparallelität bzw. der Nicht-Expertenparallelität übereinstimmt.

Referenz für das torch.sagemaker-Paket von SMP v2

Dieser Abschnitt ist eine Referenz für das von SMP v2 bereitgestellte torch.sagemaker-Paket.

torch.sagemaker.delayed_param.DelayedParamIniter

Eine API zur Anwendung von Verzögerte Parameterinitialisierung auf ein PyTorch-Modell.

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

Parameter

  • model (nn.Module) – Ein PyTorch-Modell zum Umschließen und Anwenden der verzögerten Parameterinitialisierungsfunktion von SMP v2.

  • init_method_using_config (aufrufbar) – Wenn Sie die SMP-v2-Implementierung mit Tensorparallelität oder unterstützte Transformer-Modelle von Hugging Face, die mit der SMP-Tensorparallelität kompatibel sind verwenden, behalten Sie für diesen Parameter den Standardwert None bei. Standardmäßig findet die DelayedParamIniter-API heraus, wie das angegebene Modell korrekt initialisiert wird. Für alle anderen Modelle müssen Sie eine benutzerdefinierte Parameter-Initialisierungsfunktion erstellen und sie Ihrem Skript hinzufügen. Der folgende Codeausschnitt ist die init_method_using_config-Standardfunktion, die SMP v2 für die Transformer-Modelle von Hugging Face, die mit der SMP-Tensorparallelität kompatibel sind implementiert hat. Verwenden Sie den folgenden Codeausschnitt als Referenz, um Ihre eigene Initialisierungskonfigurationsfunktion zu erstellen, sie Ihrem Skript hinzuzufügen und sie an den init_method_using_config-Parameter der SMP-DelayedParamIniter-API zu übergeben.

    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)

    Weitere Informationen zu den torch.sagemaker.module_util-Funktionen im vorherigen Codeausschnitt finden Sie unter Util-Funktionen und -Eigenschaften von torch.sagemaker.

  • verbose (Boolean) – Ob eine detailliertere Protokollierung während der Initialisierung und Validierung aktiviert werden soll. Der Standardwert ist False.

Methoden

  • get_param_init_fn() – Gibt die Parameter-Initialisierungsfunktion zurück, die Sie an das param_init_fn-Argument der PyTorch-FSDP-Wrapper-Klasse übergeben können.

  • get_post_param_init_fn() – Gibt die Parameter-Initialisierungsfunktion zurück, die Sie an das post_param_init_fn-Argument der PyTorch-FSDP-Wrapper-Klasse übergeben können. Dies ist erforderlich, wenn Sie gebundene Gewichtungen im Modell haben. Das Modell muss die Methode tie_weights implementieren. Weitere Informationen finden Sie in den Hinweisen zu gebundenen Gewichtungen in Verzögerte Parameterinitialisierung.

  • count_num_params (module: nn.Module, *args: Tuple[nn.Parameter]) – Verfolgt, wie viele Parameter von der Parameter-Initialisierungsfunktion initialisiert werden. Dies hilft bei der Implementierung der folgenden validate_params_and_buffers_inited-Methode. Normalerweise müssen Sie diese Funktion nicht explizit aufrufen, da die validate_params_and_buffers_inited-Methode diese Methode implizit im Backend aufruft.

  • validate_params_and_buffers_inited (enabled: bool=True) – Dies ist ein Kontextmanager, mit dessen Hilfe überprüft werden kann, ob die Anzahl der initialisierten Parameter mit der Gesamtzahl der Parameter im Modell übereinstimmt. Außerdem wird überprüft, ob sich alle Parameter und Puffer jetzt auf GPU-Geräten statt auf Metageräten befinden. Es werden AssertionErrors angezeigt, wenn diese Bedingungen nicht erfüllt sind. Dieser Kontextmanager ist optional und muss für die Initialisierung von Parametern nicht verwendet werden.

torch.sagemaker.distributed.checkpoint.state_dict_saver.async_save

Eingabe-API für asynchrones Speichern. Verwenden Sie diese Methode, um ein state_dict asynchron an einem bestimmten checkpoint_id zu speichern.

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:

Parameter

  • state_dict (dict) – Erforderlich. Das zu speichernde Zustandswörterbuch.

  • checkpoint_id (str) – Erforderlich. Der Speicherpfad, in dem Prüfpunkte gespeichert werden sollen.

  • storage_writer (StorageWriter) – Optional. Eine Instance von StorageWriter in PyTorch zur Ausführung von Schreiboperationen. Wenn dies nicht angegeben ist, wird die Standardkonfiguration von StorageWriter verwendet.

  • planner (SavePlanner) – Optional. Eine Instance von SavePlanner in PyTorch. Wenn dies nicht angegeben ist, wird die Standardkonfiguration von SavePlanner verwendet.

  • process_group (ProcessGroup) – Optional. Die zu bearbeitende Prozessgruppe. Falls None, wird die standardmäßige (globale) Prozessgruppe verwendet.

  • coordinator_rank (int) – Optional. Der Rang des Koordinators bei der Durchführung kollektiver Kommunikationsoperatoren wie z. B. AllReduce.

  • queue (AsyncRequestQueue) – Optional. Der zu verwendende asynchrone Scheduler. Standardmäßig wird der globale Parameter DEFAULT_ASYNC_REQUEST_QUEUE verwendet.

  • sharded_strategy (PyTorchDistSaveShardedStrategy) – Optional. Die Sharding-Strategie, die zum Speichern von Prüfpunkten verwendet werden soll. Wenn dies nicht angegeben wird, wird standardmäßig torch.sagemaker.distributed.checkpoint.state_dict_saver.PyTorchDistSaveShardedStrategy verwendet.

  • wait_error_handling (bool) – Optional. Ein Flag, das angibt, ob gewartet werden soll, bis alle Ränge die Fehlerbehandlung abgeschlossen haben. Der Standardwert ist True.

  • force_check_all_plans (bool) – Optional. Ein Flag, das festlegt, ob Pläne auch im Fall eines Cache-Treffers zwangsweise über mehrere Ränge hinweg synchronisiert werden sollen. Der Standardwert ist True.

  • s3_region (str) – Optional. Die Region, in der sich der S3-Bucket befindet. Wenn nicht angegeben, wird die Region aus der checkpoint_id abgeleitet.

  • s3client_config (S3ClientConfig) – Optional. Die Datenklasse, die konfigurierbare Parameter für den S3-Client bereitstellt. Falls nicht angegeben, wird die Standardkonfiguration von S3ClientConfig verwendet. Standardmäßig ist der part_size-Parameter auf 64 MB eingestellt.

torch.sagemaker.distributed.checkpoint.state_dict_saver.maybe_finalize_async_calls

Diese Funktion ermöglicht es einem Trainingsprozess, mehrere asynchrone Anfragen zu überwachen, die zu erledigen sind.

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

Parameter

  • blocking (bool) – Optional. Wenn True, wird gewartet, bis alle aktiven Anfragen abgeschlossen sind. Andernfalls werden nur die asynchronen Anfragen beendet, die bereits abgeschlossen wurden. Der Standardwert ist True.

  • process_group (ProcessGroup) – Optional. Die zu bearbeitende Prozessgruppe. Wenn auf None gesetzt, wird die standardmäßige (globale) Prozessgruppe verwendet.

Rückgabewerte

  • Eine Liste mit den Indizes asynchroner Aufrufe, die erfolgreich abgeschlossen wurden.

torch.sagemaker.distributed.checkpoint.state_dict_saver.save

Verwenden Sie diese Methode, um ein state_dict synchron in einer bestimmten checkpoint_id zu speichern.

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:

Parameter

  • state_dict (dict) – Erforderlich. Das zu speichernde Zustandswörterbuch.

  • checkpoint_id (str) – Erforderlich. Der Speicherpfad, in dem Prüfpunkte gespeichert werden sollen.

  • storage_writer (StorageWriter) – Optional. Eine Instance von StorageWriter in PyTorch zur Ausführung von Schreiboperationen. Wenn dies nicht angegeben ist, wird die Standardkonfiguration von StorageWriter verwendet.

  • planner (SavePlanner) – Optional. Eine Instance von SavePlanner in PyTorch. Wenn dies nicht angegeben ist, wird die Standardkonfiguration von SavePlanner verwendet.

  • process_group (ProcessGroup) – Optional. Die zu bearbeitende Prozessgruppe. Falls None, wird die standardmäßige (globale) Prozessgruppe verwendet.

  • coordinator_rank (int) – Optional. Der Rang des Koordinators bei der Durchführung kollektiver Kommunikationsoperatoren wie z. B. AllReduce.

  • wait_error_handling (bool) – Optional. Ein Flag, das angibt, ob gewartet werden soll, bis alle Ränge die Fehlerbehandlung abgeschlossen haben. Der Standardwert ist True.

  • force_check_all_plans (bool) – Optional. Ein Flag, das festlegt, ob Pläne auch im Fall eines Cache-Treffers zwangsweise über mehrere Ränge hinweg synchronisiert werden sollen. Der Standardwert ist True.

  • s3_region (str) – Optional. Die Region, in der sich der S3-Bucket befindet. Wenn nicht angegeben, wird die Region aus der checkpoint_id abgeleitet.

  • s3client_config (S3ClientConfig) – Optional. Die Datenklasse, die konfigurierbare Parameter für den S3-Client bereitstellt. Falls nicht angegeben, wird die Standardkonfiguration von S3ClientConfig verwendet. Standardmäßig ist der part_size-Parameter auf 64 MB eingestellt.

torch.sagemaker.distributed.checkpoint.state_dict_loader.load

Lädt das Zustandswörterbuch eines verteilten Modells (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:

Parameter

  • state_dict (dict) – Erforderlich. Das zu ladende state_dict.

  • checkpoint_id (str) – Erforderlich. Die ID eines Prüfpunkts. Die Bedeutung von checkpoint_id hängt vom Speicher ab. Es kann ein Pfad zu einem Ordner oder zu einer Datei sein. Es kann auch ein Schlüssel sein, wenn es sich bei dem Speicher um einen Schlüsselwertspeicher handelt.

  • storage_reader (StorageReader) – Optional. Eine Instance von StorageReader in PyTorch zur Ausführung von Leseoperationen. Falls nicht angegeben, leitet das verteilte Checkpointing automatisch den Leser anhand von checkpoint_id ab. Wenn für checkpoint_id auch None angegeben ist, wird ein Ausnahmefehler ausgelöst.

  • planner (StorageReader) – Optional. Eine Instance von LoadPlanner in PyTorch. Wenn nicht angegeben, wird die Standardkonfiguration von LoadPlanner verwendet.

  • check_keys_matched (bool) – Optional. Wenn diese Option aktiviert ist, wird geprüft, ob die state_dict-Schlüssel aller Ränge mit AllGather übereinstimmen.

  • s3_region (str) – Optional. Die Region, in der sich der S3-Bucket befindet. Wenn nicht angegeben, wird die Region aus der checkpoint_id abgeleitet.

  • s3client_config (S3ClientConfig) – Optional. Die Datenklasse, die konfigurierbare Parameter für den S3-Client bereitstellt. Falls nicht angegeben, wird die Standardkonfiguration von S3ClientConfig verwendet. Standardmäßig ist der part_size-Parameter auf 64 MB eingestellt.

torch.sagemaker.moe.moe_config.MoEConfig

Eine Konfigurationsklasse zum Einrichten der SMP-Implementierung von Mixture-of-Experts (MoE). Sie können MoE-Konfigurationswerte über diese Klasse angeben und sie an den API-Aufruf torch.sagemaker.transform übergeben. Weitere Informationen zur Verwendung dieser Klasse für das Training von MoE-Modellen finden Sie unter Expertenparallelität.

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 )

Parameter

  • smp_moe (Boolean) – Ob die SMP-Implementierung von MoE verwendet werden soll. Der Standardwert ist True.

  • random_seed (Ganzzahl) – Eine Seed-Zahl für die Zufallsoperationen in von verteilten Modulen mit Expertenparallelität. Dieser Seed-Wert wird dem expertenparallelen Rang hinzugefügt, um den tatsächlichen Seed-Wert für jeden Rang festzulegen. Er ist für jeden expertenparallelen Rang einzigartig. Der Standardwert ist 12345.

  • moe_load_balancing (Zeichenfolge) – Geben Sie den Load-Balancer-Typ des MoE-Routers an. Gültige Optionen sind aux_loss, sinkhorn, balanced und none. Der Standardwert ist sinkhorn.

  • global_token_shuffle (Boolean) – Gibt an, ob Token zwischen EP-Rängen innerhalb derselben EP-Gruppe gemischt werden sollen. Der Standardwert ist False.

  • moe_all_to_all_dispatcher (Boolean) – Gibt an, ob der All-to-All-Dispatcher für die Kommunikation im MoE verwendet werden soll. Der Standardwert ist True.

  • moe_aux_loss_coeff (Float) – Ein Koeffizient für den Hilfslastausgleichsverlust. Der Standardwert ist 0.001.

  • moe_z_loss_coeff (Float) – Koeffizient für den Z-Verlust. Der Standardwert ist 0.001.

torch.sagemaker.nn.attn.FlashSelfAttention

Eine API zur Verwendung von FlashAttention mit 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, )

Parameter

  • attention_dropout_prob (float) – Die Dropout-Wahrscheinlichkeit für die Aufmerksamkeit. Der Standardwert ist 0.0.

  • scale (float) – Falls übergeben, wird dieser Skalierungsfaktor für Softmax angewendet. Wenn auf None gesetzt (was auch der Standardwert ist), ist der Skalierungsfaktor 1 / sqrt(attention_head_size). Der Standardwert ist None.

  • triton_flash_attention (bool) — Falls übergeben, wird die Triton-Implementierung von Flash Attention verwendet. Dies ist notwendig, um Attention with Linear Biases (ALiBi) zu unterstützen (siehe den folgenden use_alibi-Parameter). Diese Version des Kernels unterstützt kein Dropout. Der Standardwert ist False.

  • use_alibi (bool) – Falls übergeben, wird Attention with Linear Biases (ALiBi) unter Verwendung der bereitgestellten Maske aktiviert. Die Verwendung von ALiBi erfordert eine Aufmerksamkeitsmaske, die wie folgt vorbereitet ist. Der Standardwert ist 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

Methoden

  • forward(self, qkv, attn_mask=None, causal=False, cast_dtype=None, layout="b h s d") – Eine reguläre PyTorch-Modulfunktion. Wenn ein module(x) aufgerufen wird, führt SMP diese Funktion automatisch aus.

    • qkv – torch.Tensor in der folgenden Form: (batch_size x seqlen x (3 x num_heads) x head_size) oder (batch_size, (3 x num_heads) x seqlen x head_size), ein Tupel von torch.Tensors, die jeweils die Form (batch_size x seqlen x num_heads x head_size) oder (batch_size x num_heads x seqlen x head_size) haben können. Basierend auf der Form muss ein geeignetes Layoutargument übergeben werden.

    • attn_mask – torch.Tensor der folgenden Form (batch_size x 1 x 1 x seqlen). Um diesen Parameter für die Aufmerksamkeitsmaske zu aktivieren, sind triton_flash_attention=True und use_alibi=True erforderlich. Informationen zum Generieren einer Aufmerksamkeitsmaske mit dieser Methode finden Sie in den Codebeispielen unter FlashAttention. Der Standardwert ist None.

    • causal – Wenn dieser Wert auf False gesetzt ist, was der Standardwert des Arguments ist, wird keine Maske angewendet. Wenn er auf True gesetzt ist, verwendet die forward-Methode die untere dreieckige Standardmaske. Der Standardwert ist False.

    • cast_dtype – Wenn dieser Wert auf einen bestimmten dtype gesetzt ist, werden die qkv-Tensoren vor attn in diesen dtype umgewandelt. Dies ist nützlich für Implementierungen wie das GPT-NeoX-Modell von Hugging Face Transformer, das nach den rotativen Einbettungen über q und k mit fp32 verfügt. Wenn auf None gesetzt, wird kein Cast angewendet. Der Standardwert ist None.

    • layout (Zeichenfolge) – Verfügbare Werte sind b h s d oder b s h d. Dies sollte auf das Layout der übergebenen qkv-Tensoren eingestellt werden, damit entsprechende Transformationen auf attn angewendet werden können. Der Standardwert ist b h s d.

Rückgabewerte

Ein einziger torch.Tensor mit der Form (batch_size x num_heads x seq_len x head_size).

torch.sagemaker.nn.attn.FlashGroupedQueryAttention

Eine API zur Verwendung von FlashGroupedQueryAttention mit SMP v2. Weitere Informationen zur Verwendung dieser API finden Sie unter Verwenden Sie FlashAttention Kernel für die Aufmerksamkeit bei gruppierten Abfragen.

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

Parameter

  • attention_dropout_prob (float) – Die Dropout-Wahrscheinlichkeit für die Aufmerksamkeit. Der Standardwert ist 0.0.

  • scale (float) – Falls übergeben, wird dieser Skalierungsfaktor für Softmax angewendet. Wenn auf None gesetzt, wird 1 / sqrt(attention_head_size) als Skalierungsfaktor verwendet. Der Standardwert ist None.

Methoden

  • forward(self, q, kv, causal=False, cast_dtype=None, layout="b s h d") – Eine reguläre PyTorch-Modulfunktion. Wenn ein module(x) aufgerufen wird, führt SMP diese Funktion automatisch aus.

    • q – torch.Tensor in der folgenden Form (batch_size x seqlen x num_heads x head_size) oder (batch_size x num_heads x seqlen x head_size). Basierend auf der Form muss ein geeignetes Layoutargument übergeben werden.

    • kv – torch.Tensor in der folgenden Form: (batch_size x seqlen x (2 x num_heads) x head_size) oder (batch_size, (2 x num_heads) x seqlen x head_size), oder ein Tupel von zwei torch.Tensoren, die jeweils die Form (batch_size x seqlen x num_heads x head_size) oder (batch_size x num_heads x seqlen x head_size) haben können. Basierend auf der Form muss auch ein entsprechendes layout-Argument übergeben werden.

    • causal – Wenn dieser Wert auf False gesetzt ist, was der Standardwert des Arguments ist, wird keine Maske angewendet. Wenn er auf True gesetzt ist, verwendet die forward-Methode die untere dreieckige Standardmaske. Der Standardwert ist False.

    • cast_dtype – Wenn dieser Wert auf einen bestimmten dtype gesetzt ist, werden die qkv-Tensoren vor attn in diesen dtype umgewandelt. Dies ist nützlich für Implementierungen wie das Hugging Face Transformer GPT-NeoX, das nach den rotativen Einbettungen über q,k mit fp32 verfügt. Wenn auf None gesetzt, wird kein Cast angewendet. Der Standardwert ist None.

    • Layout (Zeichenfolge) – Verfügbare Werte sind "b h s d" oder "b s h d". Dies sollte auf das Layout der übergebenen qkv-Tensoren eingestellt werden, damit entsprechende Transformationen auf attn angewendet werden können. Der Standardwert ist "b h s d".

Rückgabewerte

Gibt einen einzelnen torch.Tensor (batch_size x num_heads x seq_len x head_size) zurück, der das Ergebnis der Aufmerksamkeitsberechnung darstellt.

torch.sagemaker.nn.huggingface.llama_flashattn.LlamaFlashAttention

Eine API, die FlashAttention für das Llama-Modell unterstützt. Diese API verwendet die torch.sagemaker.nn.attn.FlashGroupedQueryAttention-API auf niedriger Ebene. Weitere Informationen zu ihrer Verwendung erhalten Sie unter Verwenden Sie FlashAttention Kernel für die Aufmerksamkeit bei gruppierten Abfragen.

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

Parameter

  • config – Eine FlashAttention-Konfiguration für das Llama-Modell.

Methoden

  • forward(self, hidden_states, attention_mask, position_ids, past_key_value, output_attentions, use_cache)

    • hidden_states (torch.Tensor) – Versteckte Zustände eines Tensors in Form von (batch_size x seq_len x num_heads x head_size).

    • attention_mask (torch.LongTensor) – Maske mit der Form (batch_size x seqlen), um zu vermeiden, dass Aufmerksamkeit auf das Auffüllen von Token-Indizes in Form gerichtet wird. Der Standardwert ist None.

    • position_ids (torch.LongTensor) – Wenn dieser Wert nicht auf None gesetzt ist, gibt er in der Form (batch_size x seqlen) die Positionsindizes jedes Eingabesequenz-Tokens in den Positionseinbettungen an. Der Standardwert ist None.

    • past_key_value (Cache) – Vorberechnete versteckte Zustände (Schlüssel und Werte in den Selbstaufmerksamkeitsblöcken und in den Queraufmerksamkeitsblöcken). Der Standardwert ist None.

    • output_attentions (bool) – Gibt an, ob die Aufmerksamkeitstensoren aller Aufmerksamkeitsebenen zurückgegeben werden sollen. Der Standardwert ist False.

    • use_cache (bool) – Gibt an, ob der Schlüsselwertstatus für past_key_values zurückgegeben werden sollen. Der Standardwert ist False.

Rückgabewerte

Gibt einen einzelnen torch.Tensor (batch_size x num_heads x seq_len x head_size) zurück, der das Ergebnis der Aufmerksamkeitsberechnung darstellt.

torch.sagemaker.transform

SMP v2 bietet diese torch.sagemaker.transform()-API zur Transformation von Transformer-Modellen von Hugging Face in SMP-Modellimplementierungen und zur Aktivierung der SMP-Tensorparallelität.

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, cp_comm_type: str = "p2p" )

SMP v2 verwaltet die Transformationsrichtlinien für Transformer-Modelle von Hugging Face, die mit der SMP-Tensorparallelität kompatibel sind, indem es die Konfiguration der Transformer-Modelle von Hugging Face in die SMP-Transformer-Konfiguration konvertiert.

Parameter

  • model (torch.nn.Module) – Ein Modell von Transformer-Modelle von Hugging Face, die mit der SMP-Tensorparallelität kompatibel sind zur Transformation und Anwendung des Tensorparallelitäts-Feature der SMP-Bibliothek.

  • device (torch.device) – Falls übergeben, wird auf diesem Gerät ein neues Modell erstellt. Wenn das ursprüngliche Modul einen Parameter auf dem Metagerät hat (siehe Verzögerte Parameterinitialisierung), dann wird das transformierte Modul auch auf dem Metagerät erstellt, wobei das hier übergebene Argument ignoriert wird. Der Standardwert ist None.

  • dtype (torch.dtype) – Falls übergeben, wird dieser Wert als dtype-Kontextmanager für die Erstellung des Modells festgelegt und ein Modell mit diesem dtype wird erstellt. Dies ist normalerweise unnötig, da wir das Modell bei der Verwendung von MixedPrecision mit fp32 erstellen möchten und fp32 der Standard-dtype in PyTorch ist. Der Standardwert ist None.

  • config (dict) – Dies ist ein Wörterbuch für die Konfiguration des SMP-Transformers. Der Standardwert ist None.

  • load_state_dict_from_rank0 (Boolean) – Standardmäßig erstellt dieses Modul eine neue Instance des Modells mit neuen Gewichtungen. Wenn dieses Argument auf True gesetzt ist, versucht SMP, das Zustandswörterbuch des ursprünglichen PyTorch-Modells vom 0-ten Rang in ein transformiertes Modell für die Tensorparallelgruppe zu laden, zu der der 0-te Rang gehört. Wenn dieser Wert auf True gesetzt ist, kann Rang 0 keine Parameter auf dem Metagerät haben. Nur die erste Tensorparallelgruppe füllt nach diesem Transformationsaufruf die Gewichtungen aus dem 0ten Rang auf. Sie müssen sync_module_states im FSDP-Wrapper auf True setzen, um diese Gewichtungen von der ersten Tensorparallelgruppe auf alle anderen Prozesse zu übertragen. Wenn diese Option aktiviert ist, lädt die SMP-Bibliothek das Statuswörterbuch aus dem Originalmodell. Die SMP-Bibliothek nimmt das state_dict des Modells vor der Transformation, konvertiert es, sodass es der Struktur des transformierten Modells entspricht, fragmentiert es für jeden tensorparallelen Rang, übermittelt diesen Zustand vom 0-ten Rang an andere Ränge in der Tensorparallelgruppe, zu der der 0-te Rang gehört, und lädt es. Der Standardwert ist False.

  • cp_comm_type (str) – Bestimmt die Implementierung der Kontextparallelität und ist nur anwendbar, wenn der context_parallel_degree größer als 1 ist. Verfügbare Werte für diesen Parameter sind p2p und all_gather. Die p2p-Implementierung verwendet Peer-to-Peer-Aufrufe zum Senden/Empfangen für die Ansammlung von Schlüssel- und Werte-Tensoren (KV) während der Aufmerksamkeitsberechnung. Sie läuft asynchron, sodass Überschneidungen von Kommunikation und Datenverarbeitung möglich sind. Andererseits verwendet die all_gather-Implementierung die kollektive AllGather-Kommunikationsoperation für die Ansammlung von KV-Tensoren. Der Standardwert ist "p2p".

Gibt zurück

Gibt ein transformiertes Modell zurück, das Sie mit PyTorch FSDP umschließen können. Wenn load_state_dict_from_rank0 auf True gesetzt ist, hat die Tensorparallelgruppe, die Rang 0 beinhaltet, Gewichtungen, die aus dem ursprünglichen Zustandswörterbuch auf Rang 0 geladen wurden. Bei der Anwendung von Verzögerte Parameterinitialisierung auf das ursprüngliche Modell haben nur diese Ränge die tatsächlichen Tensoren auf den CPUs für die Parameter und Puffer des transformierten Modells. Die restlichen Ränge haben weiterhin die Parameter und Puffer auf dem Metagerät, um Speicherplatz zu sparen.

Util-Funktionen und -Eigenschaften von torch.sagemaker

Util-Funktionen von torch.sagemaker
  • torch.sagemaker.init(config: Optional[Union[str, Dict[str, Any]]] = None) -> None – Initialisiert den PyTorch-Trainingsjob mit SMP.

  • torch.sagemaker.is_initialized() -> bool – Prüft, ob der Trainingsjob mit SMP initialisiert ist. Beim Rückgriff auf das native PyTorch, während der Job mit SMP initialisiert ist, sind einige der Eigenschaften nicht relevant und ändern sich in None, wie in der folgenden Eigenschaftenliste angegeben.

  • torch.sagemaker.utils.module_utils.empty_module_params(module: nn.Module, device: Optional[torch.device] = None, recurse: bool = False) -> nn.Module – Erzeugt leere Parameter für das angegebene device, sofern vorhanden, und kann für alle verschachtelten Module, falls angegeben, rekursiv sein.

  • torch.sagemaker.utils.module_utils.move_buffers_to_device(module: nn.Module, device: torch.device, recurse: bool = False) -> nn.Module – Verschiebt Modulpuffer in das angegebene device und kann, für alle verschachtelten Module, falls angegeben, rekursiv sein.

Eigenschaften

torch.sagemaker.state enthält nach der Initialisierung von SMP mit torch.sagemaker.init mehrere nützliche Eigenschaften.

  • torch.sagemaker.state.hybrid_shard_degree (int) – Der Parallelitätsgrad fragmentierter Daten gemäß der Benutzereingabe in der SMP-Konfiguration, die an torch.sagemaker.init() übergeben wurde. Weitere Informationen hierzu finden Sie unter Verwendung der SageMaker-Modellparallelitätsbibliothek v2.

  • torch.sagemaker.state.rank (int) – Der globale Rang für das Gerät im Bereich von [0, world_size).

  • torch.sagemaker.state.rep_rank_process_group (torch.distributed.ProcessGroup) – Die Prozessgruppe, die alle Geräte mit demselben Replikationsrang umfasst. Beachten Sie den subtilen, aber grundlegenden Unterschied zu torch.sagemaker.state.tp_process_group. Beim Rückgriff auf das native PyTorch wird dieser Wert auf None zurückgesetzt.

  • torch.sagemaker.state.tensor_parallel_degree (int) – Der Grad der Tensorparallelität gemäß der Benutzereingabe in der SMP-Konfiguration, die an torch.sagemaker.init() übergeben wurde. Weitere Informationen hierzu finden Sie unter Verwendung der SageMaker-Modellparallelitätsbibliothek v2.

  • torch.sagemaker.state.tp_size (int) – Ein Alias für torch.sagemaker.state.tensor_parallel_degree.

  • torch.sagemaker.state.tp_rank (int) – Der Tensorparallelitätsrang für das Gerät im Bereich von [0, tp_size), der durch den Grad der Tensorparallelität und den Rangfolgemechanismus bestimmt wird.

  • torch.sagemaker.state.tp_process_group (torch.distributed.ProcessGroup) – Die tensorparallele Prozessgruppe, die alle Geräte mit demselben Rang in anderen Dimensionen (z. B. Parallelität fragmentierter Daten und Replikation), aber einzigartigen Rängen der Tensorparallelität umfasst. Beim Rückgriff auf das native PyTorch wird dieser Wert auf None zurückgesetzt.

  • torch.sagemaker.state.world_size (int) – Die Gesamtzahl der im Training verwendeten Geräte.

Upgrade von SMP v1 auf SMP v2

Für den Wechsel von SMP v1 zu SMP v2 müssen Sie Änderungen am Skript vornehmen, um die SMP-v1-APIs zu entfernen und die SMP-v2-APIs anzuwenden. Anstatt mit Ihrem SMP-v1-Skript zu beginnen, empfehlen wir Ihnen, mit einem PyTorch-FSDP-Skript zu beginnen und den Anweisungen unter Verwendung der SageMaker-Modellparallelitätsbibliothek v2 zu folgen.

Um SMP-v1-Modelle auf SMP v2 zu übertragen, müssen Sie in SMP v1 das vollständige Modellstatuswörterbuch erfassen und die Übersetzungsfunktionen auf das Modellstatuswörterbuch anwenden, um es in das Modellprüfpunktformat von Hugging Face Transformers zu konvertieren. Dann können Sie in SMP v2, wie unter Checkpointing mit SMP beschrieben, die Modellprüfpunkte von Hugging Face Transformers laden und dann mit der Verwendung der PyTorch-Prüfpunkt-APIs mit SMP v2 fortfahren. Um SMP mit Ihrem PyTorch-FSDP-Modell zu verwenden, wechseln Sie zu SMP v2 und nehmen Sie Änderungen an Ihrem Trainingsskript vor, um PyTorch FSDP und andere aktuelle Features zu verwenden.

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

Die verfügbaren Übersetzungsfunktionen in SMP v1 finden Sie unter Support für Hugging Face Transformator-Modelle.

Anweisungen zum Speichern und Laden von Modellprüfpunkten in SMP v2 finden Sie unter Checkpointing mit SMP.