Üretken Veri Zekası

AWS Trainium'un işlevselliğini özel operatörlerle genişletme

Tarih:

Derin öğrenme (DL) hızla gelişen bir alandır ve uygulayıcılar sürekli olarak DL modellerini yeniliyor ve onları hızlandırmanın yollarını icat ediyor. Özel işleçler, geliştiricilerin PyTorch gibi mevcut makine öğrenimi (ML) çerçevelerinin işlevselliğini genişleterek DL inovasyonunun sınırlarını zorlamak için kullandıkları mekanizmalardan biridir. Genel olarak, bir Şebeke derin öğrenme modelinde bir katmanın matematiksel işlevini açıklar. A özel operatör geliştiricilerin derin öğrenme modelindeki bir katman için kendi matematiksel işlevlerini oluşturmalarına olanak tanır.

AWS Eğitimi ve AWS Çıkarımı2DL eğitimi ve çıkarımı için özel olarak oluşturulmuş , özel operatörleri (veya kısaca CustomOps'u) destekleyerek işlevselliklerini ve performanslarını genişletir. AWS Nöron, bu hızlandırıcıları destekleyen SDK, CustomOps için standart PyTorch arabirimini kullanır. Geliştiriciler, Trainium tabanlı kullanırken mevcut kodlarıyla kolayca başlayabilir Amazon EC2 Trn1 bulut sunucuları veya Inferentia2 tabanlı Amazon EC2 Inf2 bulut sunucuları. Bu gönderide, CustomOps'un avantajlarını, bunların Trainium'da etkili bir şekilde uygulanmasını ve Trainium destekli Trn1 bulut sunucularında CustomOps'u kullanmaya başlamanıza yardımcı olacak örnekleri ele alıyoruz.

Devam etmek için, aşağıdakiler gibi temel AWS hizmetlerine aşinalık Amazon Elastik Bilgi İşlem Bulutu (Amazon EC2) ima edilmektedir ve derin öğrenme, PyTorch ve C++ ile ilgili temel aşinalık yardımcı olacaktır.

PyTorch'taki özel operatörler ve avantajları

PyTorch için CustomOps, PyTorch C++ Frontend adlı sürüm 1.10'da ortaya çıktı ve C++ ile yazılmış CustomOps'u kaydetmek için kullanımı kolay bir mekanizma sağladı. CustomOps'un sağladığı avantajlardan bazıları şunlardır:

  • Performans optimizasyonu – CustomOps, belirli kullanım durumları için optimize edilebilir, bu da daha hızlı model çalıştırmalarına ve geliştirilmiş performansa yol açar.
  • Geliştirilmiş model ifade gücü – CustomOps ile, PyTorch tarafından sağlanan yerleşik işleçler kullanılarak kolayca ifade edilemeyen karmaşık hesaplamaları ifade edebilirsiniz.
  • Artırılmış modülerlik – Yeniden kullanılabilir bileşenlerin C++ kitaplıklarını oluşturarak daha karmaşık modeller oluşturmak için CustomOps'u yapı taşları olarak kullanabilirsiniz. Bu, geliştirme sürecini daha kolay ve daha modüler hale getirir ve hızlı deneyleri kolaylaştırır.
  • Daha fazla esneklik – CustomOps, yerleşik operatörlerin ötesinde operasyonlara olanak tanır; yani, standart operasyonlar kullanılarak uygulanmayan karmaşık operasyonları tanımlamanın esnek bir yolunu sağlar.

Özel operatörler için Trainium desteği

Trainium (ve AWS Inferentia2), Neuron SDK aracılığıyla yazılımda CustomOps'u destekler ve GPSIMD motorunu (General Purpose Single Instruction Multiple Data engine) kullanarak bunları donanımda hızlandırır. Bunların, DL modellerini geliştirirken ve yenilerken nasıl verimli CustomOps uygulamasını mümkün kıldığına ve daha fazla esneklik ve performans sağladığına bakalım.

Nöron SDK'sı

Neuron SDK, geliştiricilerin Trainium'da modelleri eğitmesine ve modelleri AWS Inferentia hızlandırıcılarına dağıtmasına yardımcı olur. PyTorch ve TensorFlow gibi çerçevelerle yerel olarak entegre olur, böylece Trn1 bulut sunucularında modelleri eğitmek için mevcut iş akışlarınızı ve uygulama kodunuzu kullanmaya devam edebilirsiniz.

Neuron SDK, CustomOps için standart PyTorch arabirimini kullanır. Geliştiriciler, C++'da CustomOps yazmak ve Neuron'un resmi operatör desteğini genişletmek için PyTorch'taki standart programlama arayüzünü kullanabilir. Neuron daha sonra bu CustomOps'ları aşağıdaki bölümde daha ayrıntılı olarak açıklanan GPSIMD motorunda verimli çalışacak şekilde derler. Bu, yeni deneysel CustomOps uygulamalarını kolaylaştırıyor ve bu temel donanım hakkında derin bir bilgi sahibi olmadan amaca yönelik donanım üzerinde hızlandırıyor.

Genel Amaçlı Tek Yönerge Çoklu Veri motoru

Trainium optimizasyonlarının temelinde, dört ana motora sahip tamamen bağımsız, heterojen bir bilgi işlem birimi olan NeuronCore mimarisi bulunur: tensör, vektör, skaler ve GPSIMD motoru. Skaler ve vektör motorları yüksek düzeyde paralelleştirilmiştir ve kayan nokta işlemleri için optimize edilmiştir. Tensör motoru, karma hassasiyetli hesaplamayı destekleyen, gücü optimize edilmiş, sistolik dizilime dayalıdır.

GPSIMD motoru, CustomOps'u çalıştırmak ve hızlandırmak için tasarlanmış genel amaçlı bir Tek Komutlu Çoklu Veri (SIMD) motorudur. Bu motor, düz hat C kodunu çalıştırabilen ve diğer NeuronCore-v512 motorlarının yanı sıra gömülü SRAM ve HBM belleklerine doğrudan hat içi erişime sahip sekiz adet tamamen programlanabilir 2 bit geniş genel amaçlı işlemciden oluşur. Bu yetenekler birlikte, CustomOps'un Trainium'da verimli bir şekilde çalıştırılmasına yardımcı olur.

Örneğin, verileri bellekten okuyan ve bunu yalnızca minimum sayıda ALU hesaplaması için kullanan TopK, LayerNorm veya ZeroCompression gibi işleçleri ele alalım. Normal CPU sistemleri, bu hesaplamalar için tamamen belleğe bağlıdır ve performans, verileri CPU'ya taşımak için gereken süre ile sınırlıdır. Trainium'da GP-SIMD motorları, 2 TB/sn bellek bant genişliğini sürdürebilen yüksek bant genişliğine sahip bir akış arabirimi kullanılarak çip üzerindeki önbelleklerle sıkı bir şekilde birleştirilir. Bu nedenle, bunun gibi CustomOps, Trainium'da çok hızlı çalıştırılabilir.

Pratikte Neuron SDK özel operatörleri

Bu gönderi için bir DLAMI varsayıyoruz (herhangi bir Ubuntu or Amazon Linux'u) bir EC2 Trn1 örneğini başlatmak için kullanılıyor (2x.large veya 32x.large). Gerekli tüm yazılımların, sürücülerin ve araçların DLAMI'lerde zaten kurulu olduğunu ve eğitimle çalışmaya başlamak için yalnızca Python ortamının etkinleştirilmesi gerektiğini unutmayın. Neuron'da bulunan CustomOps işlevine "Neuron CustomOps" adını veriyoruz.

C++ koduyla PyTorch entegrasyonu sürecine benzer şekilde, Neuron CustomOps, Torch C++ API'sinin NeuronCore bağlantı noktalı bir alt kümesi aracılığıyla bir operatörün C++ uygulamasını gerektirir. Operatörün C++ uygulamasına denir. çekirdek işlevive C++ API'sinin bağlantı noktası, CustomOps geliştirme ve model entegrasyonu için gereken her şeyi, özellikle de tensör ve skaler sınıfları içerir. c10 (farklı PyTorch kitaplıklarında düşük seviyeli C++ kodu için kullanılan bir ad alanı) ve ATen işleçlerinin bir alt kümesi (veya Otomatik Tensör, PyTorch'ta kullanılan çekirdek tensör işlemlerini sağlayan C++ kitaplığı).

The torch.h çekirdeğin NeuronCore bağlantı noktalı bir alt kümesine erişebilmeniz için çekirdeği tanımlarken başlığın dahil edilmesi gerekir. Pytorch C++ API'si:

#include <torch/torch.h>

Neuron CustomOps ayrıca bir şekil fonksiyonu. Şekil işlevi, çekirdek işleviyle aynı işlev imzasına sahiptir, ancak herhangi bir hesaplama gerçekleştirmez. Yalnızca çıkış tensörünün şeklini tanımlar, gerçek değerleri tanımlamaz.

Neuron CustomOps, kitaplıklar halinde gruplandırılır ve bunları sisteme kaydetmek için makrolar kullanılır. NEURON_LIBRARY şekil işlevinin içinden kapsam. İşlev, derleme zamanında ana bilgisayarda çalıştırılacak ve register.h Torchneuron kütüphanesinden başlık:

#include "torchneuron/register.h"

Son olarak, özel kitaplık, yükleme API'si çağrılarak oluşturulur. tedarik ediyorsa build_directory parametresi, kitaplık dosyası belirtilen dizinde saklanacaktır:

import torch_neuronx
from torch_neuronx.xla_impl import custom_op custom_op.load(
name=name,# this is the name for the library(i.e, 'relu')
compute_srcs=['CustomOP.cpp'],
shape_srcs=['shape.cpp'],
build_directory*=*os.getcwd()
)

CustomOp'u bir PyTorch modelinden kullanmak için kütüphaneyi çağırarak yüklemeniz yeterlidir. load_library API'yi açın ve Neuron CustomOp'u, CustomOps'un torch.ops ad alanı aracılığıyla PyTorch'ta çağrıldığı şekilde çağırın. biçim genellikle torch.ops.<library_name>.<operator_name>. Aşağıdaki koda bakın:

import torch
import torch_neuronx
from torch_neuronx.xla_impl import custom_op custom_op.load_library('/home/user/libmy_ops.so')
out_tensor = torch.ops.my_lib.my_op(in_tensor)

Unutmayın custom_op.load API, C++ kitaplığını oluştururken, custom_op.load_library API önceden oluşturulmuş bir kitaplık dosyasını yükler.

Örnek: MLP eğitiminde Neuron CustomOps

Başlamak için aşağıdaki adımları gerçekleştirin:

  1. EC2 Trn1 bulut sunucunuzu oluşturun ve başlatın. Bir DLAMI görüntüsü kullandığınızdan emin olun (ya Ubuntu or Amazon Linux'u, tüm gerekli Neuron yazılımı ile önceden yüklenmiş) ve 512 GB'lık bir kök birim boyutu belirttiğinizi.
  2. Örneğiniz çalışmaya başladıktan sonra, bulut sunucunuza SSH.
  3. Çalışan Trn1 bulut sunucunuza PyTorch Neuron'u (torch-neuronx) kurun. Talimatlar için bkz. MLP Eğitiminde Neuron Custom C++ Operatörleri.
  4. Örnek kodu şu adresten indirin: GitHub deposu.

Artık ortamınız kurulduğuna göre, basit bir çok katmanlı algılayıcı (MLP) modelinde kullanılacak Relu ileri ve geri işlevleri biçiminde tipik bir C++ CustomOp in Neuron'da uygulanmasını açıkladığımız bu gönderiye devam edin. Adımlar, AWS Nöron Belgeleri.

Depodaki örnek kod iki klasör gösterir:

  • ./customop_mlp/PyTorch – Bir CPU için derlenecek Relu kodunu içerir
  • ./customop_mlp/neuron – Trainium için derlenecek Relu kodunu içerir

Bir Neuron CustomOp geliştirin: Çekirdek işlevi

Çekirdek işlevinin (Neuron CustomOp) geliştirilmesi için ana bilgisayar veya geliştirme ortamı, bir Linux ortamında PyTorch 1.13 ve C++17 uyumlu bir derleyici çalıştırabilir. Bu, PyTorch için herhangi bir C++ işlevi geliştirmekle aynıdır ve geliştirme ortamında bulunması gereken kitaplıklar yalnızca PyTorch ve C++ için olanlardır. Aşağıdaki örnekte, özel Relu ileri ve geri işlevleriyle bir relu.cpp dosyası oluşturuyoruz:

#include <stdint.h>
#include <stdlib.h>
#include <torch/torch.h> torch::Tensor relu_forward(const torch::Tensor& t_in) {
torch::Tensor t_out = torch::zeros(t_in.sizes(), torch::kFloat);
auto t_in_acc = t_in.accessor<float, 2>();
auto t_out_acc = t_out.accessor<float, 2>();
auto shape = t_in.sizes();
for (int i = 0; i < shape[0]; i++) {
for (int j = 0; j < shape[1]; j++) {
t_out_acc[i][j] = t_in_acc[i][j] > 0.0 ? t_in_acc[i][j] : 0.0;
}
}
return t_out;
} torch::Tensor relu_backward(const torch::Tensor& t_grad, const torch::Tensor& t_in) {
torch::Tensor t_out = torch::zeros(t_in.sizes(), torch::kFloat);
auto t_in_acc = t_in.accessor<float, 2>();
auto t_grad_acc = t_grad.accessor<float, 2>();
auto t_out_acc = t_out.accessor<float, 2>();
auto shape = t_in.sizes();
for (int i = 0; i < shape[0]; i++) {
for (int j = 0; j < shape[1]; j++) {
t_out_acc[i][j] = t_in_acc[i][j] > 0.0 ? t_grad_acc[i][j] : 0.0;
}
}
return t_out;
}

Neuron için bir Neuron CustomOp geliştirirken, şu anda desteklenen özellikleri ve API'leri hesaba kattığınızdan emin olun. Daha fazla bilgi için bkz. Özel Operatörler API Referans Kılavuzu [Deneysel].

Neuron CustomOp'u oluşturun ve kaydedin: Şekil işlevi

Neuron CustomOp ve çalışma zamanı ortamı için yapılanma, eğitimin gerçekleşeceği Trn1 örneğidir ve Neuron CustomOp, bir neuronx-cc kitaplığı olarak derlenip kaydedilecek ve yüksek düzeyde optimize edilmiş GP-SIMD üzerinde çalışacak şekilde Neuron çalışma zamanı tarafından yorumlanacaktır. motor.

Neuron CustomOp'u oluşturmak ve kaydetmek için bir şekil işlevi oluşturmamız gerekir (shape.cpp) giriş ve çıkış tensörlerini tanımlayacak ve operatörleri kaydedecektir: relu_fwd_shape ve relu_bwd_shape fonksiyonlar. Aşağıdaki koda bakın:

#include <stdint.h>
#include <stdlib.h>
#include <torch/torch.h>
#include "torchneuron/register.h" torch::Tensor relu_fwd_shape(torch::Tensor t_in) {
torch::Tensor t_out = torch::zeros(t_in.sizes(), torch::kFloat);
return t_out;
} torch::Tensor relu_bwd_shape(torch::Tensor t_grad, torch::Tensor t_in) {
torch::Tensor t_out = torch::zeros(t_in.sizes(), torch::kFloat);
return t_out;
} NEURON_LIBRARY(my_ops, m) {
m.def("relu_forward", &relu_fwd_shape, "relu_forward");
m.def("relu_backward", &relu_bwd_shape, "relu_backward");
}

The relu_fwd_shape ve relu_bwd_shape fonksiyonlar çıkış tensörünün şeklini tanımlar (giriş tensörüyle aynı boyutta olacak şekilde). Daha sonra fonksiyonların kayıtlarını yapıyoruz. NEURON_LIBRARY dürbün.

içinde ./customop_ml/nöron deposu örneğinde, CustomOp'un oluşturulmasını ve kaydedilmesini çalıştırmak için basitçe load işlevini sistemden çağırarak bir build.py komut dosyamız var. torch_neuronx.xla_impl paketi:

import os
import torch_neuronx
from torch_neuronx.xla_impl import custom_op custom_op.load(
name='relu',
compute_srcs=['relu.cpp'],
shape_srcs=['shape.cpp'],
build_directory=os.getcwd()
)

içinde build_directory, bulmalıyız librelu.so kütüphane yüklenmeye hazır ve modelimizin eğitiminde kullanılıyor.

Neuron CustomOp ile MLP modelini oluşturun

Bu bölümde, Neuron CustomOp ile MLP modelini oluşturma adımlarını inceliyoruz.

Relu sınıfını tanımlayın

Bir MLP modelinin nasıl eğitileceğine ilişkin ayrıntılı bir açıklama için bkz. Çok Katmanlı Perceptron Eğitim Eğitimi.

CustomOp'u oluşturduktan sonra adında bir Python paketi oluşturuyoruz. my_ops.py, burada meşale autograd işlevinden devralan bir Relu PyTorch sınıfı tanımlıyoruz. Autograd işlevi, bir eğitim döngüsünde kullanılabilmesi için otomatik farklılaştırmayı uygular.

Önce librelu.so kütüphanesini yüklüyoruz, ardından statik metod dekoratörleri ile tanımlanan forward ve back fonksiyonları ile yeni sınıfı tanımlıyoruz. Bu şekilde, modeli tanımladığımızda yöntemler doğrudan çağrılabilir. Aşağıdaki koda bakın:

import torch
import torch_neuronx
from torch_neuronx.xla_impl import custom_op custom_op.load_library('librelu.so') class Relu(torch.autograd.Function):
@staticmethod
def forward(ctx, input):
ctx.save_for_backward(input)
return torch.ops.my_ops.relu_forward(input) @staticmethod
def backward(ctx, grad):
input, = ctx.saved_tensors
return torch.ops.my_ops.relu_backward(grad, input), None

MLP modelini inceleyin

Artık çok katmanlı algılayıcı modelimizi Neuron CustomOp'ımızla içe aktararak yazmaya hazırız. my_ops Relu sınıfını tanımladığımız paket:

import torch
import torch.nn as nn
from torch.nn import functional as F
import my_ops # Declare 3-layer MLP for MNIST dataset
class MLP(nn.Module):
def __init__(self, input_size = 28 * 28, output_size = 10, layers = [120, 84]):
super(MLP, self).__init__()
self.fc1 = nn.Linear(input_size, layers[0])
self.fc2 = nn.Linear(layers[0], layers[1])
self.fc3 = nn.Linear(layers[1], output_size) def forward(self, x):
f1 = self.fc1(x)
r1 = my_ops.Relu.apply(f1)
f2 = self.fc2(r1)
r2 = my_ops.Relu.apply(f2)
f3 = self.fc3(r2)
return torch.log_softmax(f3, dim=1)

Eğitim komut dosyasını çalıştırın

Artık modelimizi kullanarak eğitebiliriz. train.py sağlanan komut dosyası:

import os
import time
import torch
from model import MLP from torchvision.datasets import mnist
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor # XLA imports
import torch_xla.core.xla_model as xm # Global constants
EPOCHS = 4
WARMUP_STEPS = 2
BATCH_SIZE = 32 # Load MNIST train dataset
train_dataset = mnist.MNIST(root='./MNIST_DATA_train',
train=True, download=True, transform=ToTensor()) def main():
# Prepare data loader
train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE) # Fix the random number generator seeds for reproducibility
torch.manual_seed(0) # XLA: Specify XLA device (defaults to a NeuronCore on Trn1 instance)
device = 'xla' # Move model to device and declare optimizer and loss function
model = MLP().to(device)
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
loss_fn = torch.nn.NLLLoss() # Run the training loop
print('----------Training ---------------')
model.train()
for epoch in range(EPOCHS):
start = time.time()
for idx, (train_x, train_label) in enumerate(train_loader):
optimizer.zero_grad()
train_x = train_x.view(train_x.size(0), -1)
train_x = train_x.to(device)
train_label = train_label.to(device)
output = model(train_x)
loss = loss_fn(output, train_label)
loss.backward()
optimizer.step()
xm.mark_step() # XLA: collect ops and run them in XLA runtime
if idx < WARMUP_STEPS: # skip warmup iterations
start = time.time()
# Compute statistics for the last epoch
interval = idx - WARMUP_STEPS # skip warmup iterations
throughput = interval / (time.time() - start)
print("Train throughput (iter/sec): {}".format(throughput))
print("Final loss is {:0.4f}".format(loss.detach().to('cpu'))) # Save checkpoint for evaluation
os.makedirs("checkpoints", exist_ok=True)
checkpoint = {'state_dict': model.state_dict()}
# XLA: use xm.save instead of torch.save to ensure states are moved back to cpu
# This can prevent "XRT memory handle not found" at end of test.py execution
xm.save(checkpoint,'checkpoints/checkpoint.pt') print('----------End Training ---------------') if __name__ == '__main__':
main()

Modeli xla cihazına göndererek, model ve Relu özel operatörü, optimize edilmiş Trainium donanımı kullanılarak Neuron çalışma zamanı tarafından çalıştırılmak üzere derlenir.

Bu örnekte, Trainium ML hızlandırıcı çipinde bulunan donanım motorundan (GP-SIMD) yararlanan özel bir Relu operatörünün nasıl oluşturulacağını gösterdik. Sonuç, artık çıkarım için konuşlandırılabilen eğitimli bir PyTorch modelidir.

Sonuç

Modern son teknoloji model mimarileri, mühendislik personelinden (veri bilimcileri, makine öğrenimi mühendisleri, MLOps mühendisleri ve diğerleri) depolama, bilgi işlem, bellek ve hızlandırıcılar dahil olmak üzere gerçek altyapıya kadar artan sayıda kaynak gerektirir. Bu gereksinimler, derin öğrenme modellerini geliştirmenin ve devreye almanın maliyetini ve karmaşıklığını artırır. Trainium hızlandırıcıları, bulutta DL eğitimi için yüksek performanslı, düşük maliyetli bir çözüm sunar. Trainium'un kullanımı, bir derin öğrenme derleyicisi, çalışma zamanı ve PyTorch ve TensorFlow gibi popüler çerçevelere yerel olarak entegre edilmiş araçlar içeren Neuron SDK tarafından kolaylaştırılır. (Yazım sırasında Neuron SDK 2.9'un yalnızca özel işleçlerin geliştirilmesi için PyTorch'u desteklediğini unutmayın.)

Bu gönderide gösterildiği gibi, Trainium yalnızca modellerinizi yüksek performanslı ve verimli bir şekilde eğitmek için araçlar sağlamakla kalmaz, aynı zamanda operatörlerinizi hem eğitime hem de denemeye esneklik ve anlamlılık katacak şekilde özelleştirme yeteneği sunar.

Daha fazla bilgi için, bkz. GitHub repo.


Yazarlar Hakkında

Lorea Arrizabalaga Birleşik Krallık Kamu Sektörüne uyumlu bir Çözüm Mimarıdır ve müşterilerin Amazon SageMaker ile makine öğrenimi çözümleri tasarlamasına yardımcı olur. Ayrıca donanım hızlandırmaya adanmış Teknik Alan Topluluğunun bir parçasıdır ve AWS Inferentia ile AWS Trainium iş yüklerinin test edilmesi ve karşılaştırılmasına yardımcı olur.

Shruti Koparkar AWS'de Kıdemli Ürün Pazarlama Müdürüdür. Müşterilerin makine öğrenimi ihtiyaçları için Amazon EC2 hızlandırılmış bilgi işlem altyapısını keşfetmesine, değerlendirmesine ve benimsemesine yardımcı olur.

spot_img

En Son İstihbarat

spot_img

Bizimle sohbet

Merhaba! Size nasıl yardım edebilirim?