View a markdown version of this page

Esegui un processo di formazione parallela su modello SageMaker distribuito con Tensor Parallelism - Amazon SageMaker AI

Le traduzioni sono generate tramite traduzione automatica. In caso di conflitto tra il contenuto di una traduzione e la versione originale in Inglese, quest'ultima prevarrà.

Esegui un processo di formazione parallela su modello SageMaker distribuito con Tensor Parallelism

In questa sezione, imparerai:

  • Come configurare a SageMaker PyTorch ModelTrainer e l'opzione di parallelismo del SageMaker modello per utilizzare il parallelismo tensoriale.

  • Come adattare lo script di addestramento utilizzando i moduli smdistributed.modelparallel estesi per il parallelismo tensoriale.

Per saperne di più sui smdistributed.modelparallel moduli, consulta le API SageMaker Model Parallel nella documentazione di Python SDK. SageMaker

Solo parallelismo tensoriale

Quello che segue è un esempio di un opzione di addestramento distribuito per attivare il solo parallelismo tensoriale, senza il parallelismo della pipeline. Configura i smp_options dizionari mpi_options e per specificare le opzioni di formazione distribuite per. SageMaker PyTorch ModelTrainer

Nota

Le funzionalità estese di risparmio della memoria sono disponibili tramite Deep Learning Containers for PyTorch, che implementa la libreria di parallelismo dei SageMaker modelli v1.6.0 o successiva.

Configura un SageMaker PyTorch ModelTrainer

from sagemaker.train.configs import SourceCode, Compute, InputData from sagemaker.core import image_uris mpi_options = { "enabled" : True, "processes_per_host" : 8,               # 8 processes "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none " }                 smp_options = { "enabled":True, "parameters": { "pipeline_parallel_degree": 1,    # alias for "partitions" "placement_strategy": "cluster", "tensor_parallel_degree": 4,      # tp over 4 devices "ddp": True } }                # Retrieve the training image for the desired PyTorch version training_image = image_uris.retrieve( framework="pytorch", region="us-west-2", version='1.13.1', py_version='py36', instance_type='ml.p3.16xlarge', image_scope="training" ) smp_model_trainer = ModelTrainer( training_image=training_image, source_code=SourceCode(entry_script='your_training_script.py'), role=role, compute=Compute( instance_type='ml.p3.16xlarge', instance_count=1 ), sagemaker_session=sagemaker_session, distribution={ "smdistributed": {"modelparallel": smp_options}, "mpi": mpi_options }, base_job_name="SMD-MP-demo", ) smp_model_trainer.train(input_data_config=[ InputData(channel_name="training", data_source='s3://my_bucket/my_training_data/') ])
Suggerimento

Per trovare un elenco completo dei parametri perdistribution, vedi Parametri di configurazione per il parallelismo dei modelli nella documentazione di SageMaker Python SDK.

Adatta il tuo script di formazione PyTorch

Il seguente esempio di script di addestramento mostra come adattare la libreria di parallelismo del SageMaker modello a uno script di addestramento. In questo esempio, si presume che lo script abbia il nome your_training_script.py.

import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torchnet.dataset import SplitDataset from torchvision import datasets import smdistributed.modelparallel.torch as smp class Net(nn.Module):     def __init__(self):         super(Net, self).__init__()         self.conv1 = nn.Conv2d(1, 32, 3, 1)         self.conv2 = nn.Conv2d(32, 64, 3, 1)         self.fc1 = nn.Linear(9216, 128)         self.fc2 = nn.Linear(128, 10)     def forward(self, x):         x = self.conv1(x)         x = F.relu(x)         x = self.conv2(x)         x = F.relu(x)         x = F.max_pool2d(x, 2)         x = torch.flatten(x, 1)         x = self.fc1(x)         x = F.relu(x)         x = self.fc2(x)         return F.log_softmax(x, 1) def train(model, device, train_loader, optimizer):     model.train()     for batch_idx, (data, target) in enumerate(train_loader):         # smdistributed: Move input tensors to the GPU ID used by         # the current process, based on the set_device call.         data, target = data.to(device), target.to(device)         optimizer.zero_grad()         output = model(data)         loss = F.nll_loss(output, target, reduction="mean")         loss.backward()         optimizer.step() # smdistributed: Initialize the backend smp.init() # smdistributed: Set the device to the GPU ID used by the current process. # Input tensors should be transferred to this device. torch.cuda.set_device(smp.local_rank()) device = torch.device("cuda") # smdistributed: Download only on a single process per instance. # When this is not present, the file is corrupted by multiple processes trying # to download and extract at the same time if smp.local_rank() == 0:     dataset = datasets.MNIST("../data", train=True, download=False) smp.barrier() # smdistributed: Shard the dataset based on data parallel ranks if smp.dp_size() > 1:     partitions_dict = {f"{i}": 1 / smp.dp_size() for i in range(smp.dp_size())}     dataset = SplitDataset(dataset, partitions=partitions_dict)     dataset.select(f"{smp.dp_rank()}") train_loader = torch.utils.data.DataLoader(dataset, batch_size=64) # smdistributed: Enable tensor parallelism for all supported modules in the model # i.e., nn.Linear in this case. Alternatively, we can use # smp.set_tensor_parallelism(model.fc1, True) # to enable it only for model.fc1 with smp.tensor_parallelism():     model = Net() # smdistributed: Use the DistributedModel wrapper to distribute the # modules for which tensor parallelism is enabled model = smp.DistributedModel(model) optimizer = optim.AdaDelta(model.parameters(), lr=4.0) optimizer = smp.DistributedOptimizer(optimizer) train(model, device, train_loader, optimizer)

Parallelismo tensoriale combinato con parallelismo di pipeline

Di seguito è riportato un esempio di un'opzione di addestramento distribuita che abilita il parallelismo tensoriale combinato con il parallelismo della pipeline. Imposta i smp_options parametri mpi_options e per specificare le opzioni di parallelo del modello con il parallelismo tensoriale quando configuri un. SageMaker PyTorch ModelTrainer

Nota

Le funzionalità estese di risparmio della memoria sono disponibili tramite Deep Learning Containers for PyTorch, che implementa la libreria di parallelismo dei SageMaker modelli v1.6.0 o successiva.

Configura un SageMaker PyTorch ModelTrainer

mpi_options = { "enabled" : True, "processes_per_host" : 8,               # 8 processes "custom_mpi_options" : "--mca btl_vader_single_copy_mechanism none " }                 smp_options = { "enabled":True, "parameters": { "microbatches": 4, "pipeline_parallel_degree": 2,    # alias for "partitions" "placement_strategy": "cluster", "tensor_parallel_degree": 2,      # tp over 2 devices "ddp": True } }                # Retrieve the training image for the desired PyTorch version training_image = image_uris.retrieve( framework="pytorch", region="us-west-2", version='1.13.1', py_version='py36', instance_type='ml.p3.16xlarge', image_scope="training" ) smp_model_trainer = ModelTrainer( training_image=training_image, source_code=SourceCode(entry_script='your_training_script.py'), role=role, compute=Compute( instance_type='ml.p3.16xlarge', instance_count=1 ), sagemaker_session=sagemaker_session, distribution={ "smdistributed": {"modelparallel": smp_options}, "mpi": mpi_options }, base_job_name="SMD-MP-demo", ) smp_model_trainer.train(input_data_config=[ InputData(channel_name="training", data_source='s3://my_bucket/my_training_data/') ])

Adatta il tuo script PyTorch di formazione

Il seguente esempio di script di addestramento mostra come adattare la libreria di parallelismo del SageMaker modello a uno script di addestramento. Nota che lo script di addestramento ora include il decoratore smp.step:

import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torchnet.dataset import SplitDataset from torchvision import datasets import smdistributed.modelparallel.torch as smp class Net(nn.Module):     def __init__(self):         super(Net, self).__init__()         self.conv1 = nn.Conv2d(1, 32, 3, 1)         self.conv2 = nn.Conv2d(32, 64, 3, 1)         self.fc1 = nn.Linear(9216, 128)         self.fc2 = nn.Linear(128, 10)     def forward(self, x):         x = self.conv1(x)         x = F.relu(x)         x = self.conv2(x)         x = F.relu(x)         x = F.max_pool2d(x, 2)         x = torch.flatten(x, 1)         x = self.fc1(x)         x = F.relu(x)         x = self.fc2(x)         return F.log_softmax(x, 1) # smdistributed: Define smp.step. Return any tensors needed outside. @smp.step def train_step(model, data, target):     output = model(data)     loss = F.nll_loss(output, target, reduction="mean")     model.backward(loss)     return output, loss def train(model, device, train_loader, optimizer):     model.train()     for batch_idx, (data, target) in enumerate(train_loader):         # smdistributed: Move input tensors to the GPU ID used by         # the current process, based on the set_device call.         data, target = data.to(device), target.to(device)         optimizer.zero_grad()         # Return value, loss_mb is a StepOutput object         _, loss_mb = train_step(model, data, target)         # smdistributed: Average the loss across microbatches.         loss = loss_mb.reduce_mean()         optimizer.step() # smdistributed: Initialize the backend smp.init() # smdistributed: Set the device to the GPU ID used by the current process. # Input tensors should be transferred to this device. torch.cuda.set_device(smp.local_rank()) device = torch.device("cuda") # smdistributed: Download only on a single process per instance. # When this is not present, the file is corrupted by multiple processes trying # to download and extract at the same time if smp.local_rank() == 0:     dataset = datasets.MNIST("../data", train=True, download=False) smp.barrier() # smdistributed: Shard the dataset based on data parallel ranks if smp.dp_size() > 1:     partitions_dict = {f"{i}": 1 / smp.dp_size() for i in range(smp.dp_size())}     dataset = SplitDataset(dataset, partitions=partitions_dict)     dataset.select(f"{smp.dp_rank()}") # smdistributed: Set drop_last=True to ensure that batch size is always divisible # by the number of microbatches train_loader = torch.utils.data.DataLoader(dataset, batch_size=64, drop_last=True) model = Net() # smdistributed: enable tensor parallelism only for model.fc1 smp.set_tensor_parallelism(model.fc1, True) # smdistributed: Use the DistributedModel container to provide the model # to be partitioned across different ranks. For the rest of the script, # the returned DistributedModel object should be used in place of # the model provided for DistributedModel class instantiation. model = smp.DistributedModel(model) optimizer = optim.AdaDelta(model.parameters(), lr=4.0) optimizer = smp.DistributedOptimizer(optimizer) train(model, device, train_loader, optimizer)