آموزش JAX را در چندین پردازنده گرافیکی مقیاس‌پذیر کنید

۱. مقدمه

جکس در مسیر یادگیری GPU. آزمایشگاه 6: معرفی آموزش چند پردازنده گرافیکی.

در این آزمایشگاه کد، شما همان مرحله آموزش JAX را حفظ می‌کنید و کار را بین هر دو پردازنده گرافیکی روی گره خود تقسیم می‌کنید.

روش‌های مختلفی برای تقسیم آموزش بین دستگاه‌ها وجود دارد: موازی‌سازی داده‌ها ، موازی‌سازی تانسور یا مدل، و موازی‌سازی خط لوله. این آزمایشگاه کد، موازی‌سازی داده‌ها، ساده‌ترین نقطه شروع را پوشش می‌دهد. هر پردازنده گرافیکی (GPU) برش متفاوتی از دسته را دریافت می‌کند، همان مدل و همان مرحله آموزش را اجرا می‌کند و گرادیان‌ها را به یک به‌روزرسانی مشترک اضافه می‌کند.

بخش مهم این است که کد مرحله آموزش به ندرت تغییر می‌کند. شما نحوه قرارگیری آرایه‌ها روی دستگاه‌ها را تغییر می‌دهید و JAX اجرای توزیع‌شده را مدیریت می‌کند.

کاری که انجام خواهید داد

  • ایجاد یک Mesh ، یک شبکه منطقی از پردازنده‌های گرافیکی (GPU) با محورهای نامگذاری شده
  • دسته‌های آموزشی Shard و پارامترهای تکثیر شده با NamedSharding و PartitionSpec
  • محل قرارگیری حاصل را با jax.debug.visualize_array_sharding بررسی کنید
  • همان مرحله آموزش jax.jit را روی آرایه‌های خرد شده اجرا کنید و اجازه دهید JAX آن را موازی‌سازی کند.
  • برای کنترل صریح هر shard، محاسبه گرادیان را با shard_map بازنویسی کنید.
  • توان عملیاتی تک پردازنده گرافیکی را در مقابل چند پردازنده گرافیکی اندازه‌گیری کنید و اندازه دسته کلی را بررسی کنید

آنچه نیاز دارید

  • یک پروژه Google Cloud با قابلیت پرداخت صورتحساب، و اعتبار کارگاه یا رزرو شامل استفاده از GPU
  • سهمیه حداقل ۲ پردازنده گرافیکی NVIDIA L4 در منطقه انتخابی شما ( نحوه بررسی سهمیه پردازنده گرافیکی )
  • محیطی که JAX دو یا چند پردازنده گرافیکی (GPU) را می‌بیند. این آزمایشگاه کد در صورتی که فقط یکی از آنها قابل مشاهده باشد، روی سلول اول متوقف می‌شود.
  • تکمیل codelab های ۱ تا ۵، یا یک محیط JAX GPU معادل. Codelab 4 به طور خاص حافظه نهان Fashion-MNIST و حلقه آموزشی optax را که در اینجا دوباره استفاده شده است، تنظیم می‌کند.

زمان تخمینی برای تکمیل: ۶۰ دقیقه .

نحوه کار آموزش موازی داده

موازی‌سازی داده‌ها ساده‌ترین استراتژی چند پردازنده گرافیکی است. این استراتژی دسته را تقسیم می‌کند و مدل را تکثیر می‌کند. در اینجا فرآیند را مشاهده می‌کنید:

  1. پارامترهای مدل را تکثیر کنید تا هر پردازنده گرافیکی یک کپی کامل از وزن‌ها را در خود نگه دارد.
  2. دسته داده‌ها را در امتداد بُعد دسته خرد کنید و هر پردازنده گرافیکی (GPU) برش متفاوتی دریافت می‌کند.
  3. به طور مستقل روی هر پردازنده گرافیکی به جلو و عقب حرکت می‌کند و هر کدام گرادیان‌ها را روی برش محلی خود محاسبه می‌کنند.
  4. گرادیان‌ها را به طور میانگین در پردازنده‌های گرافیکی مختلف کاهش دهید تا هر کپی به‌روزرسانی یکسانی را مشاهده کند.
  5. به‌روزرسانی پارامترها به طور یکسان در هر پردازنده گرافیکی با گرادیان‌های یکسان به معنای وزن‌های جدید یکسان است.

وقتی هر پردازنده گرافیکی (GPU) محاسبات محلی کافی داشته باشد، موازی‌سازی داده‌ها می‌تواند مثال‌های per_device_batch * num_gpus را تنها با افزایش اندکی در زمان گام در مقایسه با پردازش یک پردازنده گرافیکی per_device_batch پردازش کند.

این جنبه‌ی مثبت ماجرا است. نکته‌ی مثبت این است که هر پردازنده‌ی گرافیکی باید یک کپی کامل از مدل را ذخیره کند، بنابراین موازی‌سازی داده‌ها وقتی خود مدل برای یک دستگاه خیلی بزرگ است، کمکی نمی‌کند. همچنین در هر مرحله به همگام‌سازی گرادیان بین دستگاه‌ها نیاز دارد که می‌تواند برای مدل‌های بسیار بزرگ، دسته‌های کوچک یا اتصالات کندتر به یک گلوگاه تبدیل شود.

۲. قبل از شروع

پروژه خود را انتخاب کنید

در کنسول گوگل کلود ، یک پروژه با قابلیت پرداخت فعال انتخاب یا ایجاد کنید.

پوسته ابری را باز کنید

برای شروع یک جلسه Cloud Shell ، روی Activate Cloud Shell (آیکون ترمینال در سمت راست بالای کنسول) کلیک کنید، سپس آن را به پروژه خود هدایت کنید:

gcloud config set project <YOUR_PROJECT_ID>

فراهم کردن محیط GPU

دستور زیر را در Cloud Shell اجرا کنید. Codelab 1 هر دستور را با جزئیات توضیح می‌دهد، از جمله سهمیه GPU و الزامات منطقه.

gcloud services enable \
  container.googleapis.com \
  compute.googleapis.com \
  iam.googleapis.com \
  cloudresourcemanager.googleapis.com \
  logging.googleapis.com \
  monitoring.googleapis.com

git clone https://github.com/Google-Cloud-AI/partner-ai-nvidia.git
cd partner-ai-nvidia/05-workshops/jax-on-gpu/terraform

cp terraform.tfvars.example terraform.tfvars

فایل terraform.tfvars را ویرایش کنید و project_id = " " قرار دهید. project_id = " " ، سپس کلاستر را آماده کرده و JupyterLab را مستقر کنید:

terraform init
terraform apply
$(terraform output -raw get_credentials_command)

cd ..
kubectl apply -f deploy/jupyter.yaml

terraform apply حدود ۱۲ دقیقه طول می‌کشد. پس از اتمام، منتظر Pod و LoadBalancer باشید، سپس توکن یکبار مصرف JupyterLab را از لاگ Pod بخوانید:

kubectl get pod jax-jupyter -w         # wait for Running, then Ctrl+C
kubectl get svc jax-jupyter-svc -w     # wait for EXTERNAL-IP, then Ctrl+C
kubectl logs jax-jupyter | grep -o 'token=[a-z0-9]*' | head -1

http:// :8884 را باز کنید http:// :8884 وارد کنید، توکن را جایگذاری کنید و یک دفترچه یادداشت پایتون ۳ جدید در /workspace ایجاد کنید. هر بلوک کد در این codelab به سلولی از آن دفترچه یادداشت می‌رود.

آنچه این codelab نیاز دارد را نصب کنید 

!pip install --quiet optax matplotlib

GPU را تنظیم و تأیید کنید

این کد JAX را به همراه سه عنصر اولیه شاردینگ مورد نیاز شما، یعنی Mesh ، PartitionSpec و NamedSharding ، که همگی از jax.sharding قابل مشاهده هستند، وارد می‌کند و آنچه JAX می‌تواند ببیند را چاپ می‌کند. خط assert len(gpu_devices) >= 2 دروازه این آزمایشگاه کد است: هر چیزی پس از آن بیش از یک دستگاه را فرض می‌کند، بنابراین اگر فقط یک GPU قابل مشاهده باشد، کد در اینجا متوقف می‌شود تا اجازه ندهد مراحل بعدی به روش‌های گیج‌کننده‌ای با شکست مواجه شوند.

import os
os.environ["LD_LIBRARY_PATH"] = "/usr/local/nvidia/lib64:" + os.environ.get("LD_LIBRARY_PATH", "")

import gzip
import gc
import hashlib
import shutil
import subprocess
import html
import math
import pathlib
import struct
import time
import urllib.request
import warnings
from functools import partial

from IPython.display import HTML, display
import matplotlib.pyplot as plt
import numpy as np

warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", message=".*ml_dtypes.*")
warnings.filterwarnings("ignore", message=".*JAX_PLATFORMS.*")

import jax
import jax.numpy as jnp
import optax
from jax.sharding import Mesh, PartitionSpec as P, NamedSharding

devices = jax.devices()
gpu_devices = [d for d in devices if d.platform == "gpu"]
NUM_DEVICES = len(gpu_devices)

print(f"JAX version:     {jax.__version__}")
print(f"Default backend: {jax.default_backend()}")
print(f"GPU devices:     {gpu_devices}")
print(f"GPU count:       {NUM_DEVICES}")

assert len(gpu_devices) >= 2, (
    f"This lab needs at least 2 GPUs. Found {len(gpu_devices)}. "
    f"Available devices: {devices}"
)


def block_tree(tree):
    """Wait until a PyTree of JAX arrays is ready on device."""
    return jax.block_until_ready(tree)


def drop_device_refs(*names, clear_compilation_cache=False):
    """Drop global references that may hold device buffers, then run cleanup."""
    for name in names:
        globals().pop(name, None)
    gc.collect()
    if clear_compilation_cache and hasattr(jax, "clear_caches"):
        jax.clear_caches()


def show_table(headers, rows, title=None, aligns=None):
    """Render rows as an HTML table."""
    aligns = aligns or ["left"] * len(headers)
    parts = ["<div style='font-family: system-ui; max-width: 980px;'>"]
    if title:
        parts.append(f"<h4 style='margin: 0 0 8px 0;'>{html.escape(title)}</h4>")
    parts.append("<table style='border-collapse: collapse; width: 100%; font-size: 13px;'>")
    parts.append("<thead><tr>")
    for h, a in zip(headers, aligns):
        parts.append(
            f"<th style='text-align:{a}; border-bottom:1px solid #d0d7de; padding:6px;'>"
            f"{html.escape(str(h))}</th>"
        )
    parts.append("</tr></thead><tbody>")
    for row in rows:
        parts.append("<tr>")
        for cell, a in zip(row, aligns):
            parts.append(
                f"<td style='text-align:{a}; border-bottom:1px solid #eef1f4; padding:6px;'>"
                f"{html.escape(str(cell))}</td>"
            )
        parts.append("</tr>")
    parts.append("</tbody></table></div>")
    display(HTML("".join(parts)))


def show_bars(rows, title, unit="", lower_is_better=False):
    """Render (label, value) pairs as a horizontal bar chart in HTML."""
    max_value = max(float(value) for _, value in rows) or 1.0
    color = "#1a7f37" if not lower_is_better else "#0969da"
    parts = ["<div style='font-family: Arial, sans-serif; max-width: 760px;'>"]
    parts.append(f"<h4 style='margin: 0 0 8px 0;'>{html.escape(title)}</h4>")
    for label, value in rows:
        width = max(3, 100 * float(value) / max_value)
        parts.append(
            "<div style='display:grid; grid-template-columns: 190px 1fr 130px; gap: 8px; "
            "align-items:center; margin: 6px 0;'>"
            f"<div style='font-size:13px;'>{html.escape(str(label))}</div>"
            "<div style='background:#f6f8fa; border-radius:6px; overflow:hidden; height:22px;'>"
            f"<div style='height:22px; width:{width:.1f}%; background:{color};'></div></div>"
            f"<div style='font-size:13px; font-variant-numeric: tabular-nums;'>{float(value):,.1f} {html.escape(unit)}</div>"
            "</div>"
        )
    parts.append(
        f"<div style='font-size:12px; color:#57606a;'>"
        f"{'Lower' if lower_is_better else 'Higher'} is better.</div></div>"
    )
    display(HTML("".join(parts)))

شما باید یک نسخه JAX، gpu به عنوان backend پیش‌فرض و لیستی از دو دستگاه CUDA با تعداد پردازنده گرافیکی ۲ را ببینید.

۳. بارگذاری Fashion-MNIST و تعریف یک MLP با محاسبات سنگین

این مرحله از همان مجموعه داده Fashion-MNIST مانند آزمایش قبلی استفاده می‌کند، اما با یک مدل محاسباتی سنگین، بنابراین اثر چند پردازنده گرافیکی (multi-GPU) راحت‌تر قابل مشاهده است. بارگذاری مجموعه داده و آماده‌سازی دسته‌ای سمت میزبان، یک بار و قبل از هر زمان‌بندی اتفاق می‌افتد. معیار بعدی در این آزمایشگاه کد، تنها مرحله آموزش پردازنده گرافیکی کامپایل شده را اندازه‌گیری می‌کند.

دانلود و آماده‌سازی داده‌ها 

DATA_DIR = pathlib.Path.home() / ".cache" / "jax-course" / "fashion-mnist"
DATA_DIR.mkdir(parents=True, exist_ok=True)

FILES = {
    "train-images-idx3-ubyte.gz": "8d4fb7e6c68d591d4c3dfef9ec88bf0d",
    "train-labels-idx1-ubyte.gz": "25c81989df183df01b3e8a0aad5dffbe",
}
PRIMARY_BASE_URL = "https://github.com/zalandoresearch/fashion-mnist/raw/master/data/fashion"

def md5sum(path):
    digest = hashlib.md5()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def download_if_needed(filename, expected_md5):
    path = DATA_DIR / filename
    if path.exists() and md5sum(path) == expected_md5:
        return path
    for base in [PRIMARY_BASE_URL]:
        try:
            print(f"Downloading {filename}")
            urllib.request.urlretrieve(f"{base}/{filename}", path)
            if md5sum(path) != expected_md5:
                raise ValueError("MD5 mismatch")
            return path
        except Exception:
            if path.exists():
                path.unlink()
    raise RuntimeError(f"Could not download {filename}")


def read_idx_images(path):
    with gzip.open(path, "rb") as f:
        _, n, rows, cols = struct.unpack(">IIII", f.read(16))
        return np.frombuffer(f.read(), dtype=np.uint8).reshape(n, rows, cols)


def read_idx_labels(path):
    with gzip.open(path, "rb") as f:
        _, n = struct.unpack(">II", f.read(8))
        return np.frombuffer(f.read(), dtype=np.uint8).reshape(n)


paths = {name: download_if_needed(name, cs) for name, cs in FILES.items()}
train_images = read_idx_images(paths["train-images-idx3-ubyte.gz"])
train_labels = read_idx_labels(paths["train-labels-idx1-ubyte.gz"])

# Shuffle once on the host
perm = np.random.default_rng(0).permutation(len(train_images))
x_train_all = (train_images[perm].astype(np.float32) / 255.0).reshape(len(train_images), -1)
y_train_all = train_labels[perm].astype(np.int32)

drop_device_refs("train_images", "train_labels", "perm")

مدل بلوک مشترک را تعریف کنید

این مدل یک بلوک پنهان مشترک را بارها اعمال می‌کند. استفاده مجدد از یک بلوک، محاسبات محلی را که هر پردازنده گرافیکی در هر مرحله انجام می‌دهد، بدون افزایش تعداد مقادیر گرادیانی که باید در بین پردازنده‌های گرافیکی همگام‌سازی شوند، افزایش می‌دهد.

طراحی این تقسیم‌بندی را فعال می‌کند. هزینه همگام‌سازی گرادیان با تعداد پارامترها افزایش می‌یابد، در حالی که محاسبه با میزان محاسباتی که برای هر مثال انجام می‌دهید، افزایش می‌یابد. BLOCK_REPEATS دومی را بدون دست زدن به اولی افزایش می‌دهد.

INPUT_DIM = 28 * 28
WIDTH = 1024
NUM_CLASSES = 10
BLOCK_REPEATS = 128
BLOCK_MIX = 0.10
LEARNING_RATE = 3e-4

PER_DEVICE_BATCH = 1024
GLOBAL_BATCH = PER_DEVICE_BATCH * NUM_DEVICES
NUM_TRAIN_BATCHES = 8
BENCHMARK_WARMUP = 4
BENCHMARK_STEPS = 15
BENCHMARK_REPEATS = 3


def init_params(seed=0):
    rng = np.random.default_rng(seed)

    def normal(shape, scale):
        return rng.standard_normal(shape).astype(np.float32) * scale

    return {
        "w_in": normal((INPUT_DIM, WIDTH), math.sqrt(2.0 / INPUT_DIM)),
        "b_in": np.zeros((WIDTH,), dtype=np.float32),
        "w_block": normal((WIDTH, WIDTH), math.sqrt(2.0 / WIDTH)),
        "b_block": np.zeros((WIDTH,), dtype=np.float32),
        "w_out": normal((WIDTH, NUM_CLASSES), math.sqrt(2.0 / WIDTH)),
        "b_out": np.zeros((NUM_CLASSES,), dtype=np.float32),
    }


def make_fashion_batches(batch_size, num_batches=NUM_TRAIN_BATCHES):
    needed = batch_size * num_batches
    if needed > len(x_train_all):
        raise ValueError(
            f"Need {needed:,} examples, but Fashion-MNIST has {len(x_train_all):,}."
        )
    x = x_train_all[:needed].reshape(num_batches, batch_size, INPUT_DIM)
    y = y_train_all[:needed].reshape(num_batches, batch_size)
    return x, y


def model(params, x):
    h = jax.nn.gelu(x @ params["w_in"] + params["b_in"])

    def block(h, _):
        z = jax.nn.gelu(h @ params["w_block"] + params["b_block"])
        h = (1.0 - BLOCK_MIX) * h + BLOCK_MIX * z
        return h, None

    h, _ = jax.lax.scan(block, h, xs=None, length=BLOCK_REPEATS)
    return h @ params["w_out"] + params["b_out"]


def loss_with_metrics(params, batch):
    x, y = batch
    logits = model(params, x)
    loss = optax.softmax_cross_entropy_with_integer_labels(logits, y).mean()
    accuracy = jnp.mean(jnp.argmax(logits, axis=-1) == y)
    return loss, {"accuracy": accuracy}


optimizer = optax.adamw(learning_rate=LEARNING_RATE, weight_decay=1e-4)
param_template = init_params(seed=1)
PARAM_COUNT = sum(x.size for x in param_template.values())
GRADIENT_MB = PARAM_COUNT * np.dtype(np.float32).itemsize / 1e6

drop_device_refs("param_template")

show_table(
    ["", "Value"],
    [
        ("Dataset", f"Fashion-MNIST train ({len(x_train_all):,} examples)"),
        ("Input shape", "28 x 28 grayscale, flattened to 784"),
        ("Model", f"shared-block MLP, width={WIDTH}, repeats={BLOCK_REPEATS}"),
        ("Parameters", f"{PARAM_COUNT:,}"),
        ("Gradient size", f"{GRADIENT_MB:.1f} MB per step"),
        ("Per-GPU batch", PER_DEVICE_BATCH),
        ("Global batch on all GPUs", GLOBAL_BATCH),
        ("Benchmark", f"median of {BENCHMARK_REPEATS} x {BENCHMARK_STEPS} steps"),
    ],
    title="Fashion-MNIST compute-heavy workload",
)

شما باید یک جدول خلاصه را ببینید که حجم کار را شرح می‌دهد، از جمله تعداد پارامترها، اندازه گرادیان‌هایی که در هر مرحله همگام‌سازی می‌شوند، و اندازه‌های دسته‌ای به ازای هر پردازنده گرافیکی و سراسری. در یک گره با دو پردازنده گرافیکی، دسته سراسری دو برابر دسته‌ای به ازای هر پردازنده گرافیکی است.

۴. اندازه‌گیری خط مبنای تک پردازنده گرافیکی

قبل از اضافه کردن پردازنده گرافیکی دوم، به یک عدد برای مقایسه نیاز دارید. خط پایه روی یک پردازنده گرافیکی با اندازه دسته PER_DEVICE_BATCH اجرا می‌شود - دقیقاً همان مقدار کاری که هر پردازنده گرافیکی در اجرای چند پردازنده گرافیکی انجام می‌دهد.

با پین کردن داده‌ها و پارامترها به یک دستگاه واحد شروع کنید:

single_device = gpu_devices[0]

x_batches_1gpu, y_batches_1gpu = make_fashion_batches(PER_DEVICE_BATCH)
x_batches_1gpu = jax.device_put(x_batches_1gpu, single_device)
y_batches_1gpu = jax.device_put(y_batches_1gpu, single_device)

params_1gpu = jax.device_put(init_params(seed=1), single_device)
opt_state_1gpu = optimizer.init(params_1gpu)

مرحله آموزش و معیار را تعریف کنید

train_step یک گام معمولی تک‌پردازنده‌ای با مقدار و گرادیان، به‌روزرسانی optax و پارامترهای جدید است. و دستگاه‌ها یا شاردینگ که بخشی از این آزمایشگاه کد است، وجود ندارد.

benchmark_training ابتدا گرم می‌شود تا کامپایل محاسبه نشود، سپس سه تکرار پانزده مرحله‌ای را ضرب می‌کند و میانه را گزارش می‌دهد. block_tree چیزی است که زمان‌بندی را با ارسال‌های ناهمزمان JAX منصفانه می‌کند. 

@jax.jit
def train_step(params, opt_state, batch):
    (loss, metrics), grads = jax.value_and_grad(loss_with_metrics, has_aux=True)(
        params,
        batch,
    )
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)
    return params, opt_state, {"loss": loss, "accuracy": metrics["accuracy"]}


def benchmark_training(
    step_fn,
    params,
    opt_state,
    x_batches,
    y_batches,
    warmup=BENCHMARK_WARMUP,
    steps=BENCHMARK_STEPS,
    repeats=BENCHMARK_REPEATS,
):
    """Warm up, then report median steady-state throughput."""
    num_batches = x_batches.shape[0]
    for i in range(warmup):
        batch = (x_batches[i % num_batches], y_batches[i % num_batches])
        params, opt_state, _ = step_fn(params, opt_state, batch)
    block_tree((params, opt_state))

    batch_size = x_batches.shape[1]
    timings = []
    metrics = None
    for repeat in range(repeats):
        start = time.perf_counter()
        for i in range(steps):
            batch_index = (repeat * steps + i) % num_batches
            batch = (x_batches[batch_index], y_batches[batch_index])
            params, opt_state, metrics = step_fn(params, opt_state, batch)
        params, opt_state, metrics = block_tree((params, opt_state, metrics))
        timings.append(time.perf_counter() - start)

    elapsed = float(np.median(timings))
    return {
        "examples_per_sec": steps * batch_size / elapsed,
        "ms_per_step": 1000 * elapsed / steps,
        "final_loss": float(metrics["loss"]),
        "final_accuracy": float(metrics["accuracy"]),
    }

اجرای خط پایه

دو خط آخر کد زیر، تابع drop_device_refs را فراخوانی می‌کنند که در سراسر این آزمایشگاه کد تکرار می‌شود: هر اجرا پارامترها، وضعیت بهینه‌ساز و دسته‌ها را روی GPU اختصاص می‌دهد و این بافرها تا زمانی که یک متغیر سراسری پایتون به آنها اشاره می‌کند، زنده می‌مانند. حذف نام‌ها و اجرای gc.collect() حافظه دستگاه را قبل از اینکه اجرای بعدی حافظه خود را اختصاص دهد، آزاد می‌کند، بنابراین مرحله بعدی با خطای کمبود حافظه ناشی از اجرایی که قبلاً انجام داده‌اید، مواجه نمی‌شود. 

result_1gpu = benchmark_training(
    train_step,
    params_1gpu,
    opt_state_1gpu,
    x_batches_1gpu,
    y_batches_1gpu,
)

show_table(
    ["Metric", "Value"],
    [
        ("GPUs used", "1"),
        ("Batch per step", PER_DEVICE_BATCH),
        ("Throughput", f"{result_1gpu['examples_per_sec']:,.0f} examples/sec"),
        ("Step time", f"{result_1gpu['ms_per_step']:.2f} ms"),
        ("Final loss", f"{result_1gpu['final_loss']:.4f}"),
        ("Final accuracy", f"{100 * result_1gpu['final_accuracy']:.1f}%"),
    ],
    title="Single-GPU baseline",
)

# Keep scalar timing results, but free device buffers from the single-GPU run.
drop_device_refs(
    "params_1gpu",
    "opt_state_1gpu",
    "x_batches_1gpu",
    "y_batches_1gpu",
)

این کد کمی طول می‌کشد زیرا مرحله را کامپایل می‌کند، گرم می‌شود و سپس ۴۵ مرحله زمان‌بندی شده را اجرا می‌کند. در نهایت باید یک جدول پایه تک پردازنده گرافیکی داشته باشید که توان عملیاتی را بر حسب مثال در ثانیه، زمان مرحله بر حسب میلی‌ثانیه و میزان اتلاف و دقت از آخرین مرحله را گزارش می‌دهد. result_1gpu از پاکسازی جان سالم به در می‌برد زیرا اعداد اعشاری پایتون را نگه می‌دارد، نه آرایه‌های دستگاه را.

۵. ایجاد مش دستگاه

یک Mesh ، پردازنده‌های گرافیکی فیزیکی را به یک شبکه منطقی با محورهای نامگذاری شده نگاشت می‌کند. برای موازی‌سازی داده‌ها، یک شبکه یک بعدی با تمام پردازنده‌های گرافیکی در امتداد یک محور 'data' واحد ایجاد می‌کنید.

نام محورها برای انتخاب

این آزمایشگاه کد، محور مش را 'data' می‌نامد زیرا برای موازی‌سازی داده‌ها، دسته را تکه‌تکه می‌کند. در مدل‌های بزرگ‌تر، ممکن است از نام‌هایی مانند 'model' ، 'tensor' ، 'pipeline' یا 'fsdp' برای توصیف انواع دیگر موازی‌سازی استفاده کنید. یک مش دوبعدی ممکن است از ('data', 'model') استفاده کند، که در آن یک محور دسته را تکه‌تکه می‌کند و محور دیگر وزن‌ها یا فعال‌سازی‌ها را مدل‌سازی می‌کند.

این نام‌ها هیچ معنای خاصی برای JAX ندارند. آن‌ها فقط از طریق PartitionSpec ها و مجموعه‌هایی که به آن‌ها اشاره می‌کنند، معنادار می‌شوند.

در عمل، این سه نوع داده اولیه با هم استفاده می‌شوند: PartitionSpec طرح‌بندی را توصیف می‌کند، NamedSharding آن طرح‌بندی را به شبکه‌ای از دستگاه‌ها متصل می‌کند و jax.device_put یک آرایه را به آن طرح‌بندی منتقل می‌کند. 

mesh = Mesh(np.array(gpu_devices), ("data",))

show_table(
    ["", "Value"],
    [
        ("Mesh shape", str(mesh.shape)),
        ("Axis names", str(mesh.axis_names)),
        ("Devices", ", ".join(str(d) for d in mesh.devices.flat)),
    ],
    title="Device mesh",
)

جدول باید یک محور واحد به نام data گزارش دهد که اندازه آن برابر با تعداد پردازنده گرافیکی شما باشد و هر دو دستگاه CUDA را فهرست کند.

بررسی توپولوژی پردازنده گرافیکی (GPU)

آموزش موازی داده‌ها، گرادیان‌ها را در هر مرحله به طور کامل کاهش می‌دهد، بنابراین مسیر بین دو پردازنده گرافیکی مستقیماً روی مسیر بحرانی قرار می‌گیرد. مسیرهای NVLink که nvidia-smi آنها را NV* گزارش می‌کند، برای این کار بسیار بهتر از مسیرهای PHB هستند که از طریق پل میزبان و PCIe حرکت می‌کنند. 

if shutil.which("nvidia-smi"):
    topo = subprocess.run(
        ["nvidia-smi", "topo", "-m"],
        check=False,
        text=True,
        capture_output=True,
    )
    print(topo.stdout or topo.stderr)
else:
    print("nvidia-smi is not available in this environment.")

شما باید یک ماتریس با یک سطر و یک ستون برای هر پردازنده گرافیکی (GPU) ببینید. در یک g2-standard-24 ، دو L4 از طریق PCIe به هم متصل هستند، بنابراین انتظار داشته باشید که کد GPU0 به GPU1 یک مسیر کلاس PHB را گزارش دهد. این چیزی است که این نوع ماشین به شما می‌دهد، نه یک پیکربندی اشتباه، بلکه محدودیتی در مورد میزان مقیاس آموزش موازی داده‌ها در اینجا است و هر نتیجه‌ای را که در مرحله مقایسه به دست می‌آورید، توضیح می‌دهد.

۶. داده‌ها را خرد کنید و پارامترها را تکثیر کنید

در آموزش موازی داده، دقیقاً دو جایگذاری وجود دارد:

  • داده‌ها در امتداد بُعد دسته‌ای خرد می‌شوند و هر پردازنده گرافیکی (GPU) برش متفاوتی از دسته را دریافت می‌کند.
  • پارامترها با هر پردازنده گرافیکی که یک کپی کامل را در خود نگه می‌دارد، تکثیر می‌شوند، بنابراین مسیر رو به جلو به طور یکسان اجرا می‌شود.

PartitionSpec('data', None) بُعد اول را در محور مش 'data' تقسیم می‌کند و بُعد دوم را کپی می‌کند. PartitionSpec() بدون هیچ آرگومانی، همه چیز را کپی می‌کند.

آرایه‌های دسته‌ای یک بُعد دسته‌ای مقدم دارند، زیرا make_fashion_batches تمام دسته‌های آموزشی را که روی هم انباشته شده‌اند، برمی‌گرداند. به همین دلیل است که از P(None, "data", None) استفاده می‌کنند، بنابراین بُعد ۰ را کامل می‌گذاریم، مثال‌ها را در بُعد ۱ در GPUها خرد می‌کنیم و ویژگی‌ها را تکثیر می‌کنیم. 

batch_data_sharding = NamedSharding(mesh, P("data", None))
batch_label_sharding = NamedSharding(mesh, P("data"))
all_data_sharding = NamedSharding(mesh, P(None, "data", None))
all_label_sharding = NamedSharding(mesh, P(None, "data"))
replicated = NamedSharding(mesh, P())

x_batches_multi, y_batches_multi = make_fashion_batches(GLOBAL_BATCH)
x_batches_multi = jax.device_put(x_batches_multi, all_data_sharding)
y_batches_multi = jax.device_put(y_batches_multi, all_label_sharding)

params_multi = jax.device_put(init_params(seed=1), replicated)
opt_state_multi = optimizer.init(params_multi)

print(
    f"Global batch: {GLOBAL_BATCH} examples "
    f"({PER_DEVICE_BATCH} per GPU x {NUM_DEVICES} GPUs)"
)
print(f"Training batches shape: {x_batches_multi.shape}")
print()

print("One training batch: sharded along the batch dimension")
jax.debug.visualize_array_sharding(x_batches_multi[0])

print()
print("Weight w_block: replicated on all GPUs")
jax.debug.visualize_array_sharding(params_multi["w_block"])

jax.debug.visualize_array_sharding یک شبکه متنی چاپ می‌کند که نشان می‌دهد هر پردازنده گرافیکی (GPU) کدام بخش از آرایه را در اختیار دارد. در یک گره با دو پردازنده گرافیکی (GPU)، اولین خطوط خروجی باید به صورت زیر باشند: 

Global batch: 2048 examples (1024 per GPU x 2 GPUs)
Training batches shape: (8, 2048, 784)

در زیر آن، باید دسته‌ای را ببینید که به صورت دو بلوک روی هم چیده شده، یکی برای هر پردازنده گرافیکی برچسب‌گذاری شده، و w_block به صورت یک بلوک واحد که با هر دو پردازنده گرافیکی، داده‌های خرد شده و وزن‌های تکثیر شده حاشیه‌نویسی شده است، رسم شده است.

۷. اجرای همان مرحله‌ی jitted روی آرایه‌های خرد شده

کد مرحله آموزش به هیچ وجه تغییر نمی‌کند. از همان train_step کامپایل‌شده‌ای که برای خط پایه تک پردازنده گرافیکی استفاده کردید، روی ورودی‌های خرد شده استفاده می‌کند.

وقتی JAX می‌بیند که دسته در GPUها تقسیم شده و پارامترها تکثیر شده‌اند، به طور خودکار:

  1. اجرای رو به جلو روی هر برش داده‌ی GPU
  2. گرادیان‌ها را به ازای هر شارد محاسبه می‌کند
  3. یک گرادیان کاهشی به میانگین را در سراسر پردازنده‌های گرافیکی (GPU) درج می‌کند.
  4. پارامترها را به طور یکسان در هر پردازنده گرافیکی به‌روزرسانی می‌کند

شما هیچ کد ارتباطی نمی‌نویسید. موازی‌سازی کاملاً از نحوه قرارگیری آرایه‌ها ناشی می‌شود. 

result_multi = benchmark_training(
    train_step,
    params_multi,
    opt_state_multi,
    x_batches_multi,
    y_batches_multi,
)

show_table(
    ["Metric", "Value"],
    [
        ("GPUs used", NUM_DEVICES),
        ("Global batch", GLOBAL_BATCH),
        ("Per-GPU batch", PER_DEVICE_BATCH),
        ("Throughput", f"{result_multi['examples_per_sec']:,.0f} examples/sec"),
        ("Step time", f"{result_multi['ms_per_step']:.2f} ms"),
        ("Final loss", f"{result_multi['final_loss']:.4f}"),
        ("Final accuracy", f"{100 * result_multi['final_accuracy']:.1f}%"),
    ],
    title=f"Data-parallel training on {NUM_DEVICES} GPUs",
)

شما باید جدولی به شکل جدول پایه داشته باشید که اکنون هم دسته کلی و هم دسته به ازای هر پردازنده گرافیکی را گزارش می‌دهد. از مقایسه چشمی اعداد توان عملیاتی خودداری کنید - مرحله بعدی این کار را به درستی انجام می‌دهد و نسبت تنها رقمی است که معنی‌دار است.

۸. مقایسه‌ی توان عملیاتی تک پردازنده‌ی گرافیکی و چند پردازنده‌ی گرافیکی

در این مرحله، دسته‌ی مربوط به هر پردازنده‌ی گرافیکی در هر دو اجرا یکسان است و اجرای چند پردازنده‌ی گرافیکی در هر مرحله مثال‌های بیشتری را پردازش می‌کند، زیرا هر پردازنده‌ی گرافیکی، شارد مخصوص به خود را دریافت می‌کند.

شما یک پردازنده گرافیکی (GPU) اضافه می‌کنید و همزمان حجم کار را افزایش می‌دهید، سپس می‌پرسید که آیا توان عملیاتی همچنان ادامه دارد یا خیر. این سوال با سوال «آیا یک دسته ثابت دو برابر سریع‌تر تمام می‌شود؟» متفاوت است. در اینجا منظور از سریع‌تر، توان عملیاتی بالاتر آموزش بر حسب تعداد مثال در ثانیه است. 

speed_ratio = result_multi["examples_per_sec"] / result_1gpu["examples_per_sec"]

show_table(
    ["", "1 GPU", f"{NUM_DEVICES} GPUs", "Throughput ratio"],
    [
        ("Per-GPU batch", PER_DEVICE_BATCH, PER_DEVICE_BATCH, "same"),
        ("Global batch", PER_DEVICE_BATCH, GLOBAL_BATCH, f"{NUM_DEVICES}x"),
        (
            "Examples/sec",
            f"{result_1gpu['examples_per_sec']:,.0f}",
            f"{result_multi['examples_per_sec']:,.0f}",
            f"{speed_ratio:.2f}x",
        ),
        (
            "ms/step",
            f"{result_1gpu['ms_per_step']:.2f}",
            f"{result_multi['ms_per_step']:.2f}",
            "",
        ),
    ],
    title="Throughput: same per-GPU batch",
    aligns=["left", "right", "right", "right"],
)

show_bars(
    [
        ("1 GPU", result_1gpu["examples_per_sec"]),
        (f"{NUM_DEVICES} GPUs", result_multi["examples_per_sec"]),
    ],
    "Training throughput (examples/sec)",
    "examples/s",
)

نتیجه را صادقانه بخوانید

کد بعدی به آنچه که واقعاً اندازه‌گیری کرده‌اید، وابسته است. آن را اجرا کنید و ببینید چه چیزی به دست می‌آید. 

step_ratio = result_multi["ms_per_step"] / result_1gpu["ms_per_step"]
if speed_ratio >= 1.0:
    message = (
        f"The multi-GPU run is faster for this Fashion-MNIST workload: "
        f"throughput improves by {speed_ratio:.2f}x. Each GPU still processes "
        f"{PER_DEVICE_BATCH} examples, while the global batch increases from "
        f"{PER_DEVICE_BATCH} to {GLOBAL_BATCH}. Step time changes by {step_ratio:.2f}x, "
        f"so the larger batch translates into higher examples/sec."
    )
else:
    message = (
        f"This run is still communication-bound: throughput changes by {speed_ratio:.2f}x. "
        f"Increase BLOCK_REPEATS or PER_DEVICE_BATCH to give each GPU more local work."
    )

border_color = "#1a7f37" if speed_ratio >= 1.0 else "#d1242f"
display(HTML(
    "<div style='font-family: system-ui; max-width: 900px; "
    f"border-left: 4px solid {border_color}; padding: 10px 12px; "
    "background: #f6f8fa; margin: 12px 0;'>"
    f"{html.escape(message)}"
    "</div>"
))

اگر این نسبت برابر یا بالاتر از ۱.۰ باشد، هر پردازنده گرافیکی بار کاری محلی یکسانی را اجرا می‌کند و زمان هر مرحله کمتر از زمان هر دسته افزایش می‌یابد. اگر این نسبت کمتر از ۱.۰ باشد، اجرا به صورت ارتباط-محدود است که در آن گرادیان در مسیر PHB که در بررسی توپولوژی مشاهده کردید، هزینه بیشتری نسبت به پردازنده گرافیکی اضافی دارد.

۹. با shard_map کنترل صریح را در دست بگیرید

رویکرد پیش‌فرض، اکثر بارهای کاری موازی داده را پوشش می‌دهد. با این حال، گاهی اوقات، می‌خواهید دقیقاً کنترل کنید که هر پردازنده گرافیکی چه چیزی را محاسبه می‌کند. shard_map به شما امکان می‌دهد تابعی بنویسید که روی آرایه‌های هر shard عمل کند و از collectiveهای صریح برای برقراری ارتباط بین دستگاه‌ها استفاده کند.

درون یک تابع shard_map :

  • هر پردازنده گرافیکی، شارد محلی خود مانند (1024, 784) را دریافت می‌کند.
  • in_specs نحوه برش ورودی‌ها را اعلام می‌کند.
  • out_specs نحوه‌ی مونتاژ مجدد خروجی‌ها را اعلام می‌کند.
  • jax.lax.pmean(x, 'data') میانگین x را در تمام پردازنده‌های گرافیکی در امتداد محور 'data' محاسبه می‌کند.

توجه داشته باشید که اکنون فراخوانی‌های jax.lax.pmean همان فراخوانی‌های تماماً reduce هستند که jax.jit در مرحله قبل برای شما درج کرده بود. 

@partial(
    jax.shard_map,
    mesh=mesh,
    in_specs=(P(), P("data", None), P("data",)),
    out_specs=(P(), P(), P()),
)
def compute_grads_shardmap(params, x_shard, y_shard):
    (loss, metrics), grads = jax.value_and_grad(loss_with_metrics, has_aux=True)(
        params,
        (x_shard, y_shard),
    )
    grads = jax.lax.pmean(grads, "data")
    loss = jax.lax.pmean(loss, "data")
    accuracy = jax.lax.pmean(metrics["accuracy"], "data")
    return grads, loss, accuracy


@jax.jit
def train_step_explicit(params, opt_state, batch):
    x, y = batch
    grads, loss, accuracy = compute_grads_shardmap(params, x, y)
    updates, opt_state = optimizer.update(grads, opt_state, params)
    params = optax.apply_updates(params, updates)
    return params, opt_state, {"loss": loss, "accuracy": accuracy}

به‌روزرسانی بهینه‌ساز خارج از shard_map باقی می‌ماند. گرادیان‌ها هنگام انتشار، از قبل میانگین‌گیری شده‌اند و پارامترها تکثیر می‌شوند، بنابراین هر پردازنده گرافیکی به‌روزرسانی یکسانی را اعمال می‌کند.

حالا آن را با نسخه اتوماتیک مقایسه کنید: 

params_explicit = jax.device_put(init_params(seed=1), replicated)
opt_state_explicit = optimizer.init(params_explicit)

result_explicit = benchmark_training(
    train_step_explicit,
    params_explicit,
    opt_state_explicit,
    x_batches_multi,
    y_batches_multi,
)

show_table(
    ["Approach", "Examples/sec", "ms/step"],
    [
        (
            "jit on sharded arrays",
            f"{result_multi['examples_per_sec']:,.0f}",
            f"{result_multi['ms_per_step']:.2f}",
        ),
        (
            "shard_map explicit",
            f"{result_explicit['examples_per_sec']:,.0f}",
            f"{result_explicit['ms_per_step']:.2f}",
        ),
    ],
    title="Automatic vs explicit data parallelism",
    aligns=["left", "right", "right"],
)

drop_device_refs(
    "params_multi",
    "opt_state_multi",
    "params_explicit",
    "opt_state_explicit",
    "x_batches_multi",
    "y_batches_multi",
)

شما باید دو ردیف را ببینید که محاسبه‌ی یکسانی را به دو روش بیان می‌کنند. آن‌ها را به عنوان تأییدیه‌ای در نظر بگیرید که نسخه‌ی صریح، نسخه‌ی خودکار را بازتولید می‌کند، نه به عنوان یک مسابقه - آن‌ها کار یکسانی را انجام می‌دهند و در یک مکان به پایان می‌رسند.

۱۰. اندازه دسته کلی را جابجا کنید

موازی‌سازی داده‌ها به شما امکان می‌دهد اندازه دسته کلی را با تعداد پردازنده‌های گرافیکی (GPU) مقیاس‌بندی کنید. دسته‌های بزرگتر، سربار راه‌اندازی هسته را کاهش می‌دهند و استفاده از پردازنده گرافیکی را تا جایی بهبود می‌بخشند که حافظه یا ارتباط هر دستگاه به گلوگاه تبدیل شود.

پیمایش زیر چندین اندازه دسته سراسری را در تمام پردازنده‌های گرافیکی (GPU) آزمایش می‌کند. اندازه‌هایی که به طور مساوی بر تعداد پردازنده‌های گرافیکی شما تقسیم نمی‌شوند، نادیده گرفته می‌شوند و دسته‌ای که با شکست مواجه می‌شود، به عنوان مثال با کمبود حافظه، بدون توقف حلقه گزارش می‌شود. 

BATCH_SIZES = [256, 512, 1024, 2048, 4096]

scaling_results = []

for bs in BATCH_SIZES:
    if bs % NUM_DEVICES != 0:
        print(f"Skipping global batch {bs}: not divisible by {NUM_DEVICES} GPUs.")
        continue

    try:
        x_bs, y_bs = make_fashion_batches(bs)

        x_bs = jax.device_put(x_bs, all_data_sharding)
        y_bs = jax.device_put(y_bs, all_label_sharding)

        params_bs = jax.device_put(init_params(seed=1), replicated)
        opt_bs = optimizer.init(params_bs)

        result = benchmark_training(
            train_step,
            params_bs,
            opt_bs,
            x_bs,
            y_bs,
        )

        scaling_results.append(
            {
                "batch_size": bs,
                "per_device": bs // NUM_DEVICES,
                "examples_per_sec": result["examples_per_sec"],
                "ms_per_step": result["ms_per_step"],
            }
        )

    except Exception as e:
        print(f"Batch size {bs}: {e}")

    finally:
        drop_device_refs("x_bs", "y_bs", "params_bs", "opt_bs", "result")

این طولانی‌ترین کد در حال اجرا در آزمایشگاه کد است که در آن هر اندازه دسته، کامپایل، گرم شدن و تکرارهای زمان‌بندی شده خود را آغاز می‌کند. اکنون آنچه را که اندازه‌گیری کرده‌اید، رسم کنید: 

show_table(
    ["Global batch", "Per GPU", "Examples/sec", "ms/step"],
    [
        (
            r["batch_size"],
            r["per_device"],
            f"{r['examples_per_sec']:,.0f}",
            f"{r['ms_per_step']:.2f}",
        )
        for r in scaling_results
    ],
    title=f"Batch-size scaling on {NUM_DEVICES} GPUs",
    aligns=["right", "right", "right", "right"],
)


fig, ax = plt.subplots(figsize=(8, 5))

batches = [r["batch_size"] for r in scaling_results]
throughputs = [r["examples_per_sec"] for r in scaling_results]

ax.plot(
    batches,
    throughputs,
    "o-",
    color="#0969da",
    linewidth=2,
    markersize=8,
)

ax.set_xlabel("Global batch size")
ax.set_ylabel("Examples per second")
ax.set_title(f"Throughput vs batch size — {NUM_DEVICES} GPUs data-parallel")
ax.set_xscale("log", base=2)
ax.set_xticks(batches)
ax.set_xticklabels([str(b) for b in batches])
ax.grid(True, alpha=0.25)

fig.tight_layout()
plt.show()

شما باید برای هر اندازه دسته‌ای که تکمیل شده است، یک ردیف جدول و یک نقطه روی منحنی داشته باشید. باید ببینید که با رشد دسته و مستهلک شدن سربارهای ثابت در هر مرحله، توان عملیاتی افزایش می‌یابد، سپس با اشباع شدن GPUها یا شروع تسلط بر کاهش کلی، مسطح می‌شود. اینکه این مسطح شدن کجا اتفاق می‌افتد، یکی از ویژگی‌های این مدل روی این اتصال داخلی است و این عددی است که قبل از مقیاس‌بندی به GPUهای بیشتر، ارزش دانستن دارد.

۱۱. تمیز کردن

حجم کاری Jupyter، شامل LoadBalancer و درایو دائمی را حذف کنید: 

kubectl delete -f deploy/jupyter.yaml

کلاستر، Node Pool، VPC و حساب کاربری سرویس را از بین ببرید: 

cd terraform
terraform destroy

وقتی از شما خواسته شد، yes را تایپ کنید، سپس تأیید کنید که چیزی جا نمانده است: 

gcloud container clusters list
gcloud compute instances list

هر دو باید برای این پروژه خالی باشند. اگر فقط برای این مجموعه پروژه‌ای ایجاد کرده‌اید، می‌توانید کل پروژه را از کنسول Cloud حذف کنید.

۱۲. تبریک

شما با تغییر محل قرارگیری آرایه‌ها، یک حلقه آموزشی JAX را از یک پردازنده گرافیکی به دو پردازنده گرافیکی منتقل کردید، نه با بازنویسی مرحله آموزشی.

آنچه آموخته‌اید

  • چگونه Mesh(devices, axis_names) پردازنده‌های گرافیکی فیزیکی را به یک شبکه منطقی با محورهای نامگذاری شده نگاشت می‌کند، و اینکه این نام‌ها را شما می‌توانید انتخاب کنید
  • چگونه PartitionSpec اعلام می‌کند که کدام ابعاد آرایه به کدام محورهای مش نگاشت می‌شوند - P('data', None) بُعد دسته‌ای را خرد کرده و ویژگی‌ها را تکرار می‌کند
  • چگونه NamedSharding(mesh, spec) یک مش و یک مشخصات را در یک طرح قرارگیری برای jax.device_put ترکیب می‌کند
  • چگونه jax.debug.visualize_array_sharding نشان می‌دهد که هر GPU کدام برش را نگه می‌دارد، و چرا باید آن را پس از هر تغییر مکان اجرا کرد
  • نحوه عملکرد موازی‌سازی خودکار: jax.jit روی ورودی‌های خرد شده، محاسبات all-reduce و per-shard را برای شما وارد می‌کند، بدون هیچ تغییر کدی.
  • چگونه shard_map کنترل صریحی برای هر shard ارائه می‌دهد، با استفاده از jax.lax.pmean برای میانگین‌گیری گرادیان، زمانی که نیاز به سفارشی‌سازی الگوی ارتباطی دارید
  • نحوه‌ی رفتار مقیاس‌بندی بر اساس اندازه‌ی دسته: دسته‌های سراسری بزرگ‌تر می‌توانند تا زمانی که استفاده از پردازنده‌ی گرافیکی، حافظه یا ارتباطات به گلوگاه تبدیل نشود، توان عملیاتی را بهبود بخشند.

مراحل بعدی

  • Codelab 7: آموزش یک ترانسفورماتور از ابتدا تا انتها با Flax NNX و Orbax که مکانیزم توجه از codelab 5 را با آموزش چند پردازنده گرافیکی از این codelab ترکیب می‌کند.
  • برای اینکه به هر پردازنده گرافیکی (GPU) کار محلی بیشتری بدهید، BLOCK_REPEATS یا PER_DEVICE_BATCH را افزایش دهید، سپس مرحله مقایسه را دوباره اجرا کنید و تغییر نسبت توان عملیاتی را مشاهده کنید.
  • مقیاس‌بندی مجموعه گره‌ها را به g2-standard-48 با ۴ حافظه L4 تغییر دهید — مقدار gpu_count را در terraform.tfvars gpu_count = 4 و در deploy/jupyter.yaml برابر با nvidia.com/gpu: "4" قرار دهید — و دوباره فرآیند جابجایی اندازه دسته‌ای را روی چهار دستگاه اجرا کنید.
  • یک مش دوبعدی با محور model در کنار data امتحان کنید و w_block به جای تکرار آن، در امتداد آن خرد کنید.

اسناد مرجع