একাধিক জিপিইউ জুড়ে JAX প্রশিক্ষণ স্কেল করুন

১. ভূমিকা

Jax on GPU লার্নিং পাথ। ল্যাব ৬: মাল্টি-জিপিইউ ট্রেনিং-এর পরিচিতি।

এই কোডল্যাবে আপনি সেই একই JAX ট্রেনিং ধাপটি রাখবেন এবং আপনার নোডের উভয় GPU-এর মধ্যে কাজটি ভাগ করে দেবেন।

বিভিন্ন ডিভাইসে ট্রেনিং ভাগ করার কয়েকটি উপায় আছে: ডেটা প্যারালেলিজম , টেনসর বা মডেল প্যারালেলিজম, এবং পাইপলাইন প্যারালেলিজম। এই কোডল্যাবটি ডেটা প্যারালেলিজম নিয়ে আলোচনা করে, যা শুরু করার সবচেয়ে সহজ উপায়। প্রতিটি জিপিইউ ব্যাচের একটি ভিন্ন অংশ গ্রহণ করে, একই মডেল ও একই ট্রেনিং ধাপ চালায় এবং একটি শেয়ারড আপডেটে গ্রেডিয়েন্ট সরবরাহ করে।

গুরুত্বপূর্ণ বিষয়টি হলো, ট্রেনিং স্টেপের কোডে খুব সামান্যই পরিবর্তন হয়। আপনি ডিভাইসে অ্যারেগুলো কীভাবে রাখা হবে তা পরিবর্তন করেন, এবং JAX ডিস্ট্রিবিউটেড এক্সিকিউশনটি সামলে নেয়।

আপনি যা করবেন

  • একটি Mesh তৈরি করুন, যা হলো নামযুক্ত অক্ষসহ জিপিইউ-এর একটি যৌক্তিক গ্রিড।
  • NamedSharding এবং PartitionSpec ব্যবহার করে ট্রেনিং ব্যাচগুলোকে শার্ড করুন এবং প্যারামিটারগুলো প্রতিলিপি করুন।
  • jax.debug.visualize_array_sharding ব্যবহার করে ফলাফলস্বরূপ প্রাপ্ত প্লেসমেন্টটি পরীক্ষা করুন।
  • শার্ডেড অ্যারেগুলিতে একই jax.jit ট্রেনিং ধাপটি চালান এবং JAX-কে এটি প্যারালাইজ করতে দিন।
  • প্রতিটি শার্ডের সুস্পষ্ট নিয়ন্ত্রণের জন্য shard_map ব্যবহার করে গ্রেডিয়েন্ট গণনাটি পুনরায় লিখুন।
  • একক-জিপিইউ এর সাথে একাধিক-জিপিইউ এর থ্রুপুট পরিমাপ করুন এবং গ্লোবাল ব্যাচ সাইজ সুইপ করুন।

আপনার যা যা লাগবে

  • বিলিং সক্ষম একটি গুগল ক্লাউড প্রজেক্ট, এবং ওয়ার্কশপ ক্রেডিট অথবা জিপিইউ ব্যবহারের জন্য একটি রিজার্ভেশন।
  • আপনার নির্বাচিত অঞ্চলে কমপক্ষে ২টি এনভিডিয়া এল৪ জিপিইউ-এর জন্য কোটা ( জিপিইউ কোটা কীভাবে চেক করবেন )
  • এমন একটি পরিবেশ যেখানে JAX দুই বা ততোধিক GPU দেখতে পায়। যদি কেবল একটি GPU দৃশ্যমান থাকে, তাহলে এই কোডল্যাবটি প্রথম সেলেই থেমে যায়।
  • কোডল্যাব ১ থেকে ৫ সম্পন্ন করা, অথবা একটি সমতুল্য JAX GPU পরিবেশ। বিশেষ করে কোডল্যাব ৪, Fashion-MNIST ক্যাশে এবং এখানে পুনঃব্যবহৃত optax ট্রেনিং লুপটি সেট আপ করে।

সম্পূর্ণ করতে আনুমানিক সময়: ৬০ মিনিট ।

ডেটা-প্যারালাল ট্রেনিং কীভাবে কাজ করে

ডেটা প্যারালেলিজম হলো সবচেয়ে সহজ মাল্টি-জিপিইউ কৌশল। এটি ব্যাচকে বিভক্ত করে এবং মডেলটিকে প্রতিলিপি করে। প্রক্রিয়াটি নিচে দেওয়া হলো:

  1. মডেলের প্যারামিটারগুলো এমনভাবে প্রতিলিপি করুন যাতে প্রতিটি GPU-তে ওয়েটগুলোর একটি সম্পূর্ণ অনুলিপি থাকে।
  2. ব্যাচ ডাইমেনশন বরাবর ডেটা ব্যাচটিকে শার্ড করা হয় এবং প্রতিটি জিপিইউ একটি ভিন্ন স্লাইস পায়।
  3. প্রতিটি জিপিইউ-তে স্বাধীনভাবে ফরওয়ার্ড ও ব্যাকওয়ার্ড করা হয়, যেখানে প্রতিটি তার নিজস্ব লোকাল স্লাইসে গ্রেডিয়েন্ট গণনা করে।
  4. সমস্ত জিপিইউ জুড়ে গ্রেডিয়েন্টগুলোর গড় বের করার জন্য সেগুলোকে অল-রিডিউস করুন , যাতে প্রতিটি কপি একই আপডেট পায়।
  5. প্রতিটি জিপিইউ-তে একই গ্রেডিয়েন্ট ব্যবহার করে প্যারামিটারগুলো অভিন্নভাবে আপডেট করার অর্থ হলো একই নতুন ওয়েট।

যখন প্রতিটি GPU-এর পর্যাপ্ত স্থানীয় কম্পিউট ক্ষমতা থাকে, তখন একটি GPU দ্বারা per_device_batch প্রসেস করার তুলনায় ডেটা প্যারালেলিজম per_device_batch * num_gpus উদাহরণগুলো প্রসেস করতে পারে, যার ফলে স্টেপ টাইম সামান্যই বৃদ্ধি পায়।

এটাই হলো এর সুবিধা। এর অসুবিধা হলো, প্রতিটি জিপিইউ-কে মডেলটির একটি সম্পূর্ণ কপি সংরক্ষণ করতে হয়, তাই মডেলটি যখন একটি ডিভাইসের জন্য খুব বড় হয়ে যায়, তখন ডেটা প্যারালালিজম কোনো কাজে আসে না। এছাড়াও, এর জন্য প্রতিটি ধাপে ডিভাইসগুলোর মধ্যে গ্রেডিয়েন্ট সিনক্রোনাইজেশন প্রয়োজন হয়, যা খুব বড় মডেল, ছোট ব্যাচ বা ধীরগতির ইন্টারকানেক্টের ক্ষেত্রে একটি প্রতিবন্ধকতা হয়ে উঠতে পারে।

২. শুরু করার আগে

আপনার প্রকল্প নির্বাচন করুন

গুগল ক্লাউড কনসোলে , বিলিং সক্ষম করা আছে এমন একটি প্রজেক্ট নির্বাচন করুন বা তৈরি করুন।

ওপেন ক্লাউড শেল

একটি ক্লাউড শেল সেশন শুরু করতে অ্যাক্টিভেট ক্লাউড শেল (কনসোলের উপরের ডানদিকে থাকা টার্মিনাল আইকন)-এ ক্লিক করুন, তারপর এটিকে আপনার প্রজেক্টে নির্দেশ করুন:

gcloud config set project <YOUR_PROJECT_ID>

জিপিইউ পরিবেশ প্রস্তুত করুন

ক্লাউড শেলে নিম্নলিখিতটি চালান। কোডল্যাব ১-এ জিপিইউ কোটা এবং জোনের প্রয়োজনীয়তা সহ প্রতিটি কমান্ড বিস্তারিতভাবে ব্যাখ্যা করা হয়েছে।

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 = " " তারপর project_id = " " প্রস্তুত করুন এবং JupyterLab স্থাপন করুন:

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

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

terraform apply প্রায় ১২ মিনিট সময় লাগে। এটি শেষ হলে, Pod এবং LoadBalancer-এর জন্য অপেক্ষা করুন, তারপর Pod log থেকে ওয়ান-টাইম JupyterLab টোকেনটি পড়ুন:

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 এ একটি নতুন পাইথন ৩ নোটবুক তৈরি করুন। এই কোডল্যাবের প্রতিটি কোড ব্লক সেই নোটবুকের একটি সেলে যাবে।

এই কোডল্যাবের জন্য যা যা প্রয়োজন তা ইনস্টল করুন। 

!pip install --quiet optax matplotlib

জিপিইউ সেট আপ এবং যাচাই করুন

এই কোডটি jax.sharding থেকে JAX-কে আপনার প্রয়োজনীয় তিনটি শার্ডিং প্রিমিটিভ— Mesh , PartitionSpec , এবং NamedSharding সহ ইম্পোর্ট করে এবং 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 , এবং ২টি GPU সহ দুটি CUDA ডিভাইসের একটি তালিকা দেখতে পাবেন।

৩. Fashion-MNIST লোড করুন এবং একটি গণনা-ভারী MLP সংজ্ঞায়িত করুন।

এই ধাপে আগের ল্যাবের মতোই Fashion-MNIST ডেটাসেট ব্যবহার করা হয়েছে, তবে এবার একটি কম্পিউট-হেভি মডেল ব্যবহার করা হয়েছে, যাতে মাল্টি-জিপিইউ প্রভাব আরও সহজে দেখা যায়। ডেটাসেট লোডিং এবং হোস্ট-সাইড ব্যাচ প্রস্তুতি যেকোনো টাইমিংয়ের আগে একবারই সম্পন্ন হয়। এই কোডল্যাবের পরবর্তী বেঞ্চমার্কটি শুধুমাত্র কম্পাইল করা জিপিইউ ট্রেনিং ধাপটি পরিমাপ করে।

ডেটা ডাউনলোড এবং প্রস্তুত করুন 

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 নামের একটিমাত্র অক্ষ থাকবে, যার আকার আপনার GPU সংখ্যার সমান হবে এবং এতে উভয় CUDA ডিভাইসের তালিকা থাকবে।

জিপিইউ টপোলজি পরীক্ষা করুন

ডেটা-প্যারালাল ট্রেনিং প্রতিটি ধাপে গ্রেডিয়েন্টগুলোকে সম্পূর্ণরূপে হ্রাস করে, তাই দুটি জিপিইউ-এর মধ্যবর্তী পথটি সরাসরি ক্রিটিক্যাল পাথে অবস্থান করে। NVLink পাথ (যাকে NV* হিসেবে nvidia-smi রিপোর্ট করে) এই কাজের জন্য PHB পাথের চেয়ে অনেক ভালো, কারণ পিএইচবি পাথ হোস্ট ব্রিজ এবং পিসিআইই-এর মধ্য দিয়ে যায়। 

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.")

আপনি প্রতিটি জিপিইউ-এর জন্য একটি সারি এবং একটি কলাম সহ একটি ম্যাট্রিক্স দেখতে পাবেন। একটি g2-standard-24 এ দুটি L4 PCIe-এর মাধ্যমে সংযুক্ত থাকে, তাই আশা করা যায় যে GPU0 থেকে GPU1 কোডটি একটি PHB ক্লাস পাথ রিপোর্ট করবে। এই মেশিন টাইপটি আপনাকে এটাই দেয়, এটি কোনো ভুল কনফিগারেশন নয়, বরং এটি এখানে ডেটা-প্যারালাল ট্রেনিং কতদূর পর্যন্ত স্কেল করতে পারে তার উপর একটি সীমাবদ্ধতা, এবং এটি তুলনা ধাপে আপনার প্রাপ্ত যেকোনো ফলাফলকে ব্যাখ্যা করে।

৬. ডেটা শার্ড করুন এবং প্যারামিটারগুলো প্রতিলিপি করুন।

ডেটা-প্যারালাল ট্রেনিং-এ ঠিক দুটি প্লেসমেন্ট রয়েছে:

  • ডেটা ব্যাচ ডাইমেনশন বরাবর শার্ড করা হয় এবং প্রতিটি জিপিইউ ব্যাচের একটি ভিন্ন স্লাইস পায়।
  • প্যারামিটারগুলো প্রতিলিপি করা হয় এবং প্রতিটি জিপিইউ-তে এর একটি সম্পূর্ণ অনুলিপি থাকে, ফলে ফরোয়ার্ড পাসটি হুবহু একইভাবে চলে।

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 একটি একক ব্লক হিসাবে দেখতে পাবেন যা উভয় জিপিইউ, শার্ডেড ডেটা এবং রেপ্লিকেটেড ওয়েট দ্বারা টীকাযুক্ত থাকবে।

৭. শার্ডেড অ্যারেগুলিতে একই জিটেড ধাপটি চালান

ট্রেনিং স্টেপ কোডে কোনো পরিবর্তন হয় না। এটি সেই একই কম্পাইল করা train_step ব্যবহার করে যা আপনি সিঙ্গেল-জিপিইউ বেসলাইনের জন্য ব্যবহার করেছিলেন এবং যা শার্ডেড ইনপুটেও কাজ করে।

যখন JAX দেখে যে ব্যাচটি একাধিক GPU-তে শার্ড করা হয়েছে এবং প্যারামিটারগুলো প্রতিলিপি করা হয়েছে, তখন এটি স্বয়ংক্রিয়ভাবে:

  1. প্রতিটি জিপিইউ-এর ডেটা স্লাইসে ফরোয়ার্ড পাস চালায়।
  2. প্রতি-শার্ড গ্রেডিয়েন্ট গণনা করে
  3. জিপিইউ জুড়ে গ্রেডিয়েন্টের গড় করার জন্য একটি অল-রিডিউস সন্নিবেশ করে।
  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",
)

আপনি বেসলাইন টেবিলের মতো আকৃতির একটি টেবিল পাবেন, যেখানে এখন একটি গ্লোবাল ব্যাচ এবং একটি প্রতি-জিপিইউ ব্যাচ উভয়ই রিপোর্ট করা থাকবে। থ্রুপুটের সংখ্যাগুলো শুধু চোখে দেখে তুলনা করা থেকে বিরত থাকুন — পরবর্তী ধাপে এটি সঠিকভাবে করা হয়, এবং অনুপাতটিই একমাত্র সংখ্যা যা অর্থবহ।

৮. একক-জিপিইউ এবং একাধিক-জিপিইউ থ্রুপুটের তুলনা করুন

এই পর্যায়ে, উভয় রানেই প্রতি-জিপিইউ ব্যাচ একই থাকে এবং মাল্টি-জিপিইউ রান প্রতি ধাপে আরও বেশি উদাহরণ প্রসেস করে, কারণ প্রতিটি জিপিইউ তার নিজস্ব শার্ড পায়।

আপনি একটি জিপিইউ যোগ করুন এবং একই সাথে ওয়ার্কলোডও বাড়ান, তারপর জিজ্ঞাসা করুন থ্রুপুট তার সাথে তাল মিলিয়ে চলতে পারে কিনা। এটি "একটি নির্দিষ্ট ব্যাচ কি দ্বিগুণ দ্রুত শেষ হয়" এই প্রশ্নের মতো নয়। এখানে দ্রুততর বলতে প্রতি সেকেন্ডে উদাহরণের হিসাবে উচ্চতর ট্রেনিং থ্রুপুটকে বোঝানো হচ্ছে। 

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

যদি অনুপাতটি ১.০ বা তার বেশি হয়, তাহলে প্রতিটি GPU একই স্থানীয় ওয়ার্কলোড চালায় এবং স্টেপ টাইম ব্যাচের তুলনায় কম বৃদ্ধি পায়। যদি অনুপাতটি ১.০-এর কম হয়, তাহলে রানটি কমিউনিকেশন-বাউন্ড হয়, যেখানে টপোলজি চেকে দেখা PHB পাথের উপর গ্রেডিয়েন্ট অল-রিডিউস করতে অতিরিক্ত GPU-এর চেয়ে বেশি খরচ হয়।

৯. shard_map ব্যবহার করে সুস্পষ্ট নিয়ন্ত্রণ নিন।

ডিফল্ট পদ্ধতিটি বেশিরভাগ ডেটা-প্যারালাল ওয়ার্কলোড কভার করে। তবে, কখনও কখনও আপনি প্রতিটি জিপিইউ ঠিক কী গণনা করবে তা নিয়ন্ত্রণ করতে চান। shard_map আপনাকে এমন একটি ফাংশন লিখতে দেয় যা প্রতি-শার্ড অ্যারেতে কাজ করে এবং ডিভাইসগুলোর মধ্যে যোগাযোগের জন্য সুস্পষ্ট কালেক্টিভ ব্যবহার করে।

shard_map ফাংশনের ভিতরে:

  • প্রতিটি GPU তার স্থানীয় শার্ড পায়, যেমন (1024, 784)
  • in_specs ঘোষণা করে যে ইনপুটগুলি কীভাবে ভাগ করা হবে।
  • out_specs ঘোষণা করে যে আউটপুটগুলি কীভাবে পুনরায় একত্রিত করা হবে।
  • jax.lax.pmean(x, 'data') 'data' অক্ষ বরাবর সমস্ত GPU জুড়ে x গড় বের করে।

লক্ষ্য করুন যে এখন jax.lax.pmean কলগুলো হলো সেই অল-রিডিউস (all-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 বাইরে থাকে। গ্রেডিয়েন্টগুলো তৈরি হওয়ার সময়েই সেগুলোর গড় করা হয়ে যায় এবং প্যারামিটারগুলোও প্রতিলিপিত হয়, ফলে প্রতিটি GPU একই আপডেট প্রয়োগ করে।

এখন এটিকে স্বয়ংক্রিয় সংস্করণের সাথে তুলনা করে দেখুন: 

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

আপনি একই গণনাকে দুটি ভিন্ন উপায়ে বর্ণনা করা দুটি সারি দেখতে পাবেন। এগুলোকে একটি প্রতিযোগিতা হিসেবে না দেখে, বরং সুস্পষ্ট সংস্করণটি যে স্বয়ংক্রিয় সংস্করণটিরই অনুরূপ, তার নিশ্চিতকরণ হিসেবে বিবেচনা করুন — কারণ তারা একই কাজ করে এবং একই জায়গায় গিয়ে শেষ হয়।

১০. গ্লোবাল ব্যাচ সাইজ সুইপ করুন

ডেটা প্যারালেলিজম আপনাকে জিপিইউ-এর সংখ্যার সাথে গ্লোবাল ব্যাচ সাইজ সমন্বয় করতে দেয়। বড় ব্যাচগুলো কার্নেল লঞ্চের ওভারহেড পুষিয়ে দেয় এবং জিপিইউ-এর ব্যবহার উন্নত করে, যতক্ষণ না প্রতি-ডিভাইস মেমরি বা কমিউনিকেশন একটি বাধা হয়ে দাঁড়ায়।

নিচের এই যাচাইকরণ প্রক্রিয়াটি সমস্ত জিপিইউ জুড়ে বিভিন্ন গ্লোবাল ব্যাচ সাইজ পরীক্ষা করে। যে সাইজগুলো আপনার জিপিইউ সংখ্যা দ্বারা নিঃশেষে বিভাজ্য নয়, সেগুলো বাদ দেওয়া হয়, এবং কোনো ব্যাচ ব্যর্থ হলে, যেমন মেমোরি শেষ হয়ে যাওয়ার কারণে, লুপটি বন্ধ না করেই তা রিপোর্ট করা হয়। 

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

সম্পন্ন হওয়া প্রতিটি ব্যাচ সাইজের জন্য আপনার একটি টেবিল রো এবং কার্ভের উপর একটি পয়েন্ট পাওয়া উচিত। আপনি দেখবেন যে ব্যাচ বড় হওয়ার সাথে সাথে থ্রুপুট বাড়ে এবং প্রতি-ধাপের নির্দিষ্ট ওভারহেডগুলো পরিশোধিত হয়ে যায়, তারপর যখন জিপিইউগুলো স্যাচুরেটেড হয়ে যায় বা অল-রিডিউস প্রাধান্য পেতে শুরু করে, তখন তা স্থিতিশীল হয়ে যায়। এই স্থিতিশীলতা ঠিক কোথায় ঘটে, তা এই ইন্টারকানেক্টে এই মডেলের একটি বৈশিষ্ট্য, এবং আরও জিপিইউতে স্কেল করার আগে এই সংখ্যাটি জেনে রাখা জরুরি।

১১. পরিষ্কার করুন

লোডব্যালেন্সার এবং পার্সিস্টেন্ট ভলিউম সহ জুপিটার ওয়ার্কলোডটি মুছে ফেলুন: 

kubectl delete -f deploy/jupyter.yaml

ক্লাস্টার, নোড পুল, ভিপিসি এবং সার্ভিস অ্যাকাউন্ট ধ্বংস করুন: 

cd terraform
terraform destroy

নির্দেশিত হলে yes টাইপ করুন, তারপর নিশ্চিত করুন যে পিছনে কিছু ফেলে রাখা হয়নি: 

gcloud container clusters list
gcloud compute instances list

এই প্রজেক্টের জন্য উভয়ই খালি থাকা উচিত। যদি আপনি শুধু এই সিরিজের জন্য একটি প্রজেক্ট তৈরি করে থাকেন, তাহলে আপনি এর পরিবর্তে ক্লাউড কনসোল থেকে পুরো প্রজেক্টটি মুছে ফেলতে পারেন।

১২. অভিনন্দন

আপনি অ্যারেগুলোর অবস্থান পরিবর্তন করে একটি JAX ট্রেনিং লুপকে একটি GPU থেকে দুটি GPU-তে সরিয়েছেন, ট্রেনিং ধাপটি নতুন করে লিখে নয়।

আপনি যা শিখেছেন

  • Mesh(devices, axis_names) কীভাবে ফিজিক্যাল GPU-গুলোকে নামযুক্ত অক্ষসহ একটি লজিক্যাল গ্রিডে ম্যাপ করে, এবং সেই নামগুলো আপনিই বেছে নিতে পারেন।
  • PartitionSpec যেভাবে ঘোষণা করে কোন অ্যারে ডাইমেনশন কোন মেশ অ্যাক্সিসের সাথে ম্যাপ করবে — P('data', None) ব্যাচ ডাইমেনশনকে শার্ড করে এবং ফিচারগুলোকে রেপ্লিকেট করে।
  • কীভাবে NamedSharding(mesh, spec) একটি মেশ এবং একটি স্পেককে একত্রিত করে jax.device_put এর জন্য একটি প্লেসমেন্ট প্ল্যান তৈরি করে
  • jax.debug.visualize_array_sharding কীভাবে দেখায় কোন GPU কোন স্লাইস ধারণ করে, এবং প্রতিটি প্লেসমেন্ট পরিবর্তনের পরে কেন এটি চালানো উচিত।
  • স্বয়ংক্রিয় প্যারালালিজম যেভাবে কাজ করে: শার্ডেড ইনপুটের ক্ষেত্রে jax.jit কোনো কোড পরিবর্তন ছাড়াই আপনার হয়ে অল-রিডিউস এবং প্রতি-শার্ড গণনা সন্নিবেশ করে দেয়।
  • যখন কমিউনিকেশন প্যাটার্ন কাস্টমাইজ করার প্রয়োজন হয়, তখন গ্রেডিয়েন্ট অ্যাভারেজিংয়ের জন্য jax.lax.pmean ব্যবহার করে কীভাবে shard_map প্রতিটি শার্ডের উপর সুস্পষ্ট নিয়ন্ত্রণ প্রদান করে।
  • ব্যাচ-সাইজ স্কেলিং যেভাবে কাজ করে: বড় গ্লোবাল ব্যাচ থ্রুপুট উন্নত করতে পারে, যতক্ষণ না জিপিইউ ব্যবহার, মেমরি বা যোগাযোগ বাধা হয়ে দাঁড়ায়।

পরবর্তী পদক্ষেপ

  • কোডল্যাব ৭: Flax NNX এবং Orbax ব্যবহার করে একটি ট্রান্সফর্মারকে এন্ড-টু-এন্ড প্রশিক্ষণ দিন, যা কোডল্যাব ৫-এর অ্যাটেনশন মেকানিজমের সাথে এই কোডল্যাবের মাল্টি-জিপিইউ প্রশিক্ষণকে একত্রিত করে।
  • প্রতিটি GPU-কে আরও বেশি স্থানীয় কাজ দেওয়ার জন্য BLOCK_REPEATS বা PER_DEVICE_BATCH বাড়িয়ে দিন, তারপর তুলনা ধাপটি পুনরায় চালান এবং থ্রুপুট অনুপাতের পরিবর্তন লক্ষ্য করুন।
  • ৪টি L4 ব্যবহার করে নোড পুলকে g2-standard-48 স্কেলে উন্নীত করুন — terraform.tfvars এ gpu_count = 4 এবং deploy/jupyter.yaml এ nvidia.com/gpu: "4" সেট করুন — এবং চারটি ডিভাইসে ব্যাচ-সাইজ সুইপটি পুনরায় চালান।
  • data পাশাপাশি একটি model অক্ষ সহ একটি দ্বি-মাত্রিক মেশ ব্যবহার করে দেখুন, এবং w_block কে প্রতিলিপি করার পরিবর্তে সেই অক্ষ বরাবর শার্ড করুন।

রেফারেন্স নথি