Flax NNX এবং Orbax ব্যবহার করে একটি ট্রান্সফর্মারকে প্রান্ত থেকে প্রান্তে প্রশিক্ষণ দিন।

১. ভূমিকা

এই কোডল্যাবে, আপনি JAX-এ Flax NNX ব্যবহার করে একটি ছোট ডিকোডার ট্রান্সফরমার সংজ্ঞায়িত করবেন, আপনার নোডের উভয় GPU জুড়ে শেক্সপিয়রের টেক্সটের উপর এটিকে প্রশিক্ষণ দেবেন, Orbax দিয়ে এটি সংরক্ষণ ও পুনরুদ্ধার করবেন এবং প্রশিক্ষিত ওয়েটগুলো থেকে নতুন টেক্সট তৈরি করবেন।

মডেলটি ইচ্ছাকৃতভাবেই ক্ষুদ্র রাখা হয়েছে, যাতে রয়েছে ৪টি লেয়ার, ২৫৬-মাত্রিক এমবেডিং এবং একটি বাইট-স্তরের শব্দভাণ্ডার, ফলে এটি দুটি L4 GPU-তে এক মিনিটেরও কম সময়ে প্রশিক্ষিত হয়। এর আর্কিটেকচার এবং প্রশিক্ষণের ধরণগুলো অনেক বড় মডেলে ব্যবহৃত পদ্ধতির মতোই।

আপনি যা করবেন

• Flax NNX ব্যবহার করে nnx.Embed , nnx.MultiHeadAttention , nnx.Linear এবং nnx.LayerNorm এর সাহায্যে একটি ডিকোডার ট্রান্সফরমার সংজ্ঞায়িত করুন।
• অ্যাটেনশন ব্যাকএন্ড হিসেবে কজাল jax.nn.dot_product_attention যুক্ত করুন।
• প্রথমে একটি জিপিইউ-তে এবং তারপর সবগুলো জিপিইউ-তে nnx.Optimizer এবং Optax AdamW ব্যবহার করে বাইট-লেভেল TinyShakespeare-এ প্রশিক্ষণ দিন।
• টোকেন/সেকেন্ড এককে থ্রুপুট পরিমাপ করুন এবং দুটি রানের তুলনা করুন।
• Orbax StandardCheckpointer ব্যবহার করে মডেল প্যারামিটার সংরক্ষণ এবং পুনরুদ্ধার করুন।
• প্রশিক্ষিত মডেল থেকে শেক্সপিয়রের মতো লেখা তৈরি করুন

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

• বিলিং সক্ষম একটি গুগল ক্লাউড প্রজেক্ট, এবং ওয়ার্কশপ ক্রেডিট অথবা জিপিইউ ব্যবহারের জন্য একটি রিজার্ভেশন।
• আপনার নির্বাচিত অঞ্চলে কমপক্ষে ২টি এনভিডিয়া এল৪ জিপিইউ-এর জন্য কোটা ( জিপিইউ কোটা কীভাবে চেক করবেন )
• কোডল্যাব ১ থেকে ৬ সম্পন্ন, অথবা কমপক্ষে দুটি জিপিইউ সহ একটি সমতুল্য JAX জিপিইউ পরিবেশ।
• পড থেকে বহির্গামী ইন্টারনেট সংযোগ রয়েছে, যাতে প্রথমবার চালানোর সময় টাইনিশেক্সপিয়ার টেক্সট ফাইলটি ডাউনলোড করা যায়।

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

আপনি যে স্থাপত্য নির্মাণ করছেন

মডেলটি হলো ট্রান্সফরমার ব্লকের একটি স্তূপ। প্রতিটি ব্লকের দুটি উপ-স্তর রয়েছে, যার প্রতিটি একটি রেসিডুয়াল সংযোগে আবৃত:

1. আত্ম-মনোযোগ — প্রতিটি অবস্থান তার পূর্ববর্তী সকল অবস্থানের প্রতি মনোযোগ দেয় (কার্যকারণ মুখোশ)।
2. ফিড-ফরোয়ার্ড নেটওয়ার্ক (FFN) — GELU অ্যাক্টিভেশন সহ দুটি লিনিয়ার লেয়ার, যা রিপ্রেজেন্টেশনকে প্রথমে প্রসারিত এবং পরে সংকুচিত করে।

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

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

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

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

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

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

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 flax optax orbax-checkpoint matplotlib

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

JAX, Flax NNX, Optax, এবং Orbax ইম্পোর্ট করুন, এবং কন্টেইনারটি কয়টি GPU দেখতে পাচ্ছে তা পরীক্ষা করুন।

import os

os.environ["LD_LIBRARY_PATH"] = "/usr/local/nvidia/lib64:" + os.environ.get(
    "LD_LIBRARY_PATH", ""
)
import hashlib
import html
import math
import pathlib
import time
import urllib.request
import warnings
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 flax import nnx
import orbax.checkpoint as ocp
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 lesson needs at least 2 GPUs. Found {len(gpu_devices)}. "
    f"Available devices: {devices}"
)

def block_tree(tree):
    return jax.block_until_ready(tree)

def show_table(headers, rows, title=None, aligns=None):
    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):
    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):,.0f} {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 , দুটি CUDA ডিভাইসের একটি তালিকা এবং GPU count: 2 দেখতে পাবেন। block_tree , show_table এবং show_bars হেল্পারগুলো পরবর্তী ধাপগুলোতে ব্যবহৃত টেবিল ও বার চার্ট রেন্ডার করে।

৩. বাইট-স্তরের টাইনিশেক্সপিয়ার ডেটা প্রস্তুত করুন

টাইনিশেক্সপিয়ার হলো প্রায় ১ মেগাবাইটের একটি একক টেক্সট ফাইল, যেখানে শেক্সপিয়ারের বিভিন্ন রচনা একসাথে জুড়ে দেওয়া হয়েছে। এই কোডল্যাবটি বাইট-লেভেল টোকেনাইজেশন ব্যবহার করে, যেখানে UTF-8 টেক্সটের প্রতিটি বাইট একটি টোকেনে পরিণত হয়। এটি সম্ভাব্য মানের সংখ্যা ২৫৬-তে সীমাবদ্ধ করে এবং যেকোনো টোকেনাইজার নির্ভরতা দূর করে।

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

ফাইলটি ডাউনলোড করতে, এর চেকসাম যাচাই করতে, এটিকে ট্রেন ও ভ্যালিডেশন সিকোয়েন্সে বিভক্ত করতে এবং ট্রেনিং সেটটি শাফেল করতে কোডটি চালান:

SHAKESPEARE_URL = "https://raw.githubusercontent.com/karpathy/char-rnn/master/data/tinyshakespeare/input.txt"
SHAKESPEARE_MD5 = "d015dc5942f9b2908e24d4827a3e7a5e"
DATA_DIR = pathlib.Path.home() / ".cache" / "jax-course"
DATA_DIR.mkdir(parents=True, exist_ok=True)
DATA_FILE = DATA_DIR / "tinyshakespeare.txt"

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

if DATA_FILE.exists() and md5sum(DATA_FILE) == SHAKESPEARE_MD5:
    print("Using cached tinyshakespeare.txt")
else:
    print(f"Downloading tinyshakespeare.txt")
    urllib.request.urlretrieve(SHAKESPEARE_URL, DATA_FILE)
raw_text = DATA_FILE.read_text()
data = np.frombuffer(raw_text.encode("utf-8"), dtype=np.uint8).astype(np.int32)
print()
VOCAB_SIZE = 256
SEQ_LEN = 256
PER_DEVICE_BATCH = 32
num_sequences = len(data) // SEQ_LEN
data = data[: num_sequences * SEQ_LEN].reshape(num_sequences, SEQ_LEN)
num_train = int(0.9 * num_sequences)
train_data = data[:num_train]
val_data = data[num_train:]
rng = np.random.default_rng(0)
train_data = train_data[rng.permutation(num_train)]

def make_batches(data, batch_size):
    usable = (len(data) // batch_size) * batch_size
    return data[:usable].reshape(-1, batch_size, SEQ_LEN)

show_table(
    ["", "Value"],
    [
        ("Total bytes", f"{len(raw_text):,}"),
        ("Vocabulary", f"{VOCAB_SIZE} (raw bytes)"),
        ("Sequence length", SEQ_LEN),
        ("Training sequences", f"{num_train:,}"),
        ("Validation sequences", f"{len(val_data):,}"),
    ],
    title="TinyShakespeare — byte-level tokenization",
)
print()
print("Sample text (first 200 bytes):")
print(raw_text[:200])

আপনি ‘ Using cached tinyshakespeare.txt অথবা একটি ডাউনলোড বার্তা দেখতে পাবেন। এর পরে ডেটাসেটের গঠন সম্পর্কিত একটি সারণী থাকবে, যেখানে মোট বাইট, ২৫৬টি কাঁচা বাইটের শব্দভাণ্ডার, সিকোয়েন্সের দৈর্ঘ্য, এবং প্রশিক্ষণ ও যাচাইকরণ সিকোয়েন্সের সংখ্যা উল্লেখ থাকবে। সবশেষে টেক্সটটির প্রথম ২০০ বাইট থাকবে, যাতে আপনি দেখতে পারেন মডেলটি কী থেকে শিখছে।

এখানে সেট করা তিনটি ধ্রুবক কোডল্যাবের বাকি অংশের জন্য গুরুত্বপূর্ণ। VOCAB_SIZE হলো ২৫৬, কারণ একটি বাইটের ২৫৬টি সম্ভাব্য মান থাকে, এবং SEQ_LEN হলো প্রতিটি প্রশিক্ষণ উদাহরণের জন্য ২৫৬টি অবস্থান।

PER_DEVICE_BATCH মান ৩২ এবং এটি সিঙ্গেল-জিপিইউ ও মাল্টি-জিপিইউ উভয় রানের জন্যই স্থির থাকে। প্রতি-জিপিইউ ব্যাচের এই স্থির মানটিই পরবর্তী থ্রুপুট তুলনাকে একটি দুর্বল-স্কেলিং তুলনা করে তোলে।

৪. Flax NNX দিয়ে ট্রান্সফরমারটি সংজ্ঞায়িত করুন

Flax NNX হলো JAX-এ নিউরাল নেটওয়ার্কের জন্য একটি সরলীকৃত API। আপনি লেয়ারগুলোকে পাইথন অবজেক্ট হিসেবে সংজ্ঞায়িত করেন, যেগুলো তাদের ওয়েট ইনিশিয়ালাইজেশন এবং ফরোয়ার্ড পাসের দায়িত্বে থাকে। এই ধাপে ট্রান্সফর্মারের প্রয়োজনীয় NNX-এর প্রতিটি অংশের সাথে পরিচয় করিয়ে দেওয়া হয়, তাই NNX বিষয়ে পূর্ব অভিজ্ঞতার প্রয়োজন নেই।

মডেলটি চারটি বিল্ডিং ব্লক ব্যবহার করে:

• nnx.Embed হলো একটি লুকআপ টেবিল যা একটি টোকেন ইনডেক্সকে একটি ভেক্টরের সাথে ম্যাপ করে।
• nnx.Linear হলো একটি ডেন্স ম্যাট্রিক্স মাল্টিপ্লাই এবং এর সাথে একটি ঐচ্ছিক বায়াস যুক্ত থাকে।
• nnx.LayerNorm অ্যাটেনশন এবং FFN সাব-লেয়ারের আগে ফিচারগুলোকে নর্মালাইজ করে।
• nnx.MultiHeadAttention Q/K/V প্রজেকশন, অ্যাটেনশন এবং আউটপুট প্রজেকশন পরিচালনা করে।

কার্যকারণ মনোযোগ সংযুক্ত করুন

nnx.MultiHeadAttention (B, T, D_MODEL) আকারের হিডেন স্টেট গ্রহণ করে, অভ্যন্তরীণভাবে Q, K, এবং V তৈরি করে এবং সেগুলোকে হেড-এ বিভক্ত করে। attention_fn হুকটি শুধুমাত্র মূল অ্যাটেনশন অপারেশনটি নিয়ন্ত্রণ করে, যা ঐ প্রজেকশনগুলোর পরে চলে।

NNX, attention_fn ফাংশনে একটি ড্রপআউট rng, একটি dtype, এবং একটি precision-এর মতো Flax-স্টাইলের ঐচ্ছিক আর্গুমেন্ট পাস করে। jax.nn.dot_product_attention এগুলো গ্রহণ করে না, তাই নিচের র‍্যাপারটি একটি ক্যাচ-অল ব্যবহার করে এগুলো গ্রহণ করে এবং JAX ফাংশনের জন্য শুধু প্রয়োজনীয় অংশটুকুই ফরওয়ার্ড করে।

D_MODEL = 256
NUM_HEADS = 4
FFN_DIM = 1024
NUM_LAYERS = 4
MAX_SEQ_LEN = 256
LR = 3e-4
WEIGHT_DECAY = 1e-4

def causal_sdpa(query, key, value, **_):
    return jax.nn.dot_product_attention(query, key, value, is_causal=True)

ব্লক এবং মডেল সংজ্ঞায়িত করুন

এখানে দুটি ক্লাস আছে। TransformerBlock হলো একটি অ্যাটেনশন সাব-লেয়ার ও একটি FFN সাব-লেয়ার নিয়ে গঠিত, এবং TinyTransformer এমবেডিং ও LM হেডের মধ্যে এগুলোর num_layers স্ট্যাক করে। প্রতিটি ক্লাস nnx.Module এক্সটেন্ড করে এবং এর সমস্ত লেয়ার __init__ ফাংশনে তৈরি করে।

class TransformerBlock(nnx.Module):
    def __init__(self, d_model: int, num_heads: int, ffn_dim: int, rngs: nnx.Rngs):
        self.ln1 = nnx.LayerNorm(d_model, rngs=rngs)
        self.attn = nnx.MultiHeadAttention(
            num_heads=num_heads,
            in_features=d_model,
            decode=False,
            attention_fn=causal_sdpa,
            rngs=rngs,
        )
        self.ln2 = nnx.LayerNorm(d_model, rngs=rngs)
        self.fc_up = nnx.Linear(d_model, ffn_dim, rngs=rngs)
        self.fc_down = nnx.Linear(ffn_dim, d_model, rngs=rngs)

    def __call__(self, x):
        x = x + self.attn(self.ln1(x))
        h = jax.nn.gelu(self.fc_up(self.ln2(x)))
        x = x + self.fc_down(h)
        return x

প্রি-নর্ম ব্যবস্থাটি __call__ -এ দৃশ্যমান। x = x + self.attn(self.ln1(x)) অ্যাটেনশনের আগে নর্মালাইজ করে এবং ফলাফলটি রেসিড্যুয়াল স্ট্রিমে ফিরিয়ে দেয়, এবং FFN ব্রাঞ্চটি self.ln2 সাথে একই কাজ করে।

class TinyTransformer(nnx.Module):
    def __init__(
        self,
        vocab_size: int,
        d_model: int,
        num_heads: int,
        ffn_dim: int,
        num_layers: int,
        max_seq_len: int,
        rngs: nnx.Rngs,
    ):
        self.token_embed = nnx.Embed(vocab_size, d_model, rngs=rngs)
        self.pos_embed = nnx.Embed(max_seq_len, d_model, rngs=rngs)
        self.blocks = nnx.List(
            [
                TransformerBlock(d_model, num_heads, ffn_dim, rngs=rngs)
                for _ in range(num_layers)
            ]
        )
        self.final_norm = nnx.LayerNorm(d_model, rngs=rngs)
        self.lm_head = nnx.Linear(d_model, vocab_size, use_bias=False, rngs=rngs)

    def __call__(self, tokens):
        B, T = tokens.shape
        x = self.token_embed(tokens) + self.pos_embed(jnp.arange(T))
        for block in self.blocks:
            x = block(x)
        x = self.final_norm(x)
        return self.lm_head(x)

TinyTransformer.__call__ টোকেন এমবেডিং এবং পজিশন এমবেডিং যোগ করে, ব্লকগুলোকে ক্রমানুসারে চালায়, একটি চূড়ান্ত LayerNorm প্রয়োগ করে এবং vocab_size লজিটে প্রজেক্ট করে।

মডেলটি ইনস্ট্যানশিয়েট করুন এবং পরিদর্শন করুন

মডেলটি তৈরি করতে একটি মাত্র কলই যথেষ্ট। প্যারামিটার ইনিশিয়ালাইজেশনের জন্য প্রয়োজনীয় সমস্ত র‍্যান্ডম স্টেট nnx.Rngs সামলে নেয়। একবার এটি তৈরি হয়ে গেলে, আপনি এর প্যারামিটারগুলো গণনা করতে এবং এর মধ্য দিয়ে একটি ফরোয়ার্ড পাস চালাতে পারবেন।

model = TinyTransformer(
    VOCAB_SIZE,
    D_MODEL,
    NUM_HEADS,
    FFN_DIM,
    NUM_LAYERS,
    MAX_SEQ_LEN,
    rngs=nnx.Rngs(0),
)
param_count = sum(x.size for x in jax.tree.leaves(nnx.state(model, nnx.Param)))
show_table(
    ["", "Value"],
    [
        ("Architecture", f"Decoder-only transformer"),
        ("Layers", NUM_LAYERS),
        ("Model dimension", D_MODEL),
        ("Attention heads", f"{NUM_HEADS} (head dim = {D_MODEL // NUM_HEADS})"),
        ("FFN dimension", FFN_DIM),
        ("Vocabulary", f"{VOCAB_SIZE} (byte-level)"),
        ("Max sequence length", MAX_SEQ_LEN),
        ("Parameters", f"{param_count:,}"),
    ],
    title="TinyTransformer",
)
logits = model(jnp.zeros((1, 16), dtype=jnp.int32))
print(f"Test forward pass: input (1, 16) \u2192 logits {logits.shape}")

আপনি আর্কিটেকচারটির বর্ণনা সম্বলিত একটি টেবিল দেখতে পাবেন: ৪টি লেয়ার, মডেল ডাইমেনশন ২৫৬, হেড ডাইমেনশন ৬৪ সহ ৪টি অ্যাটেনশন হেড, একটি বাইট-লেভেল ভোকাবুলারি এবং প্যারামিটার সংখ্যা। শেষ লাইনটি টেস্ট ফরোয়ার্ড পাসের ফলাফল জানায়, যেখানে (1, 16) আকারের একটি ইনপুট থেকে (1, 16, 256) আকারের লগিট তৈরি হয় — অর্থাৎ, ১৬টি ইনপুট পজিশনের প্রতিটির জন্য ২৫৬টি বাইট ভ্যালুর উপর একটি করে ডিস্ট্রিবিউশন।

৫. NNX প্রশিক্ষণের ধাপটি লিখুন।

এখানে আপনি NNX মডিউলের অবস্থা স্বয়ংক্রিয়ভাবে পরিচালনা করতে @nnx.jit ব্যবহার করেন। এটি JIT কম্পাইলেশনের জন্য মডিউলগুলোকে স্ট্রাকচার ও অ্যারেতে বিভক্ত করে, তারপর আপডেট করা অ্যারেগুলোকে আবার একত্রিত করে। ফলে, আপনি ধাপটি এমনভাবে লেখেন যেন মডিউলগুলো সাধারণ পাইথন অবজেক্ট।

লস হলো পরবর্তী টোকেনের পূর্বাভাস। প্রতিটি অবস্থানে মডেলটি পরবর্তী টোকেনটির পূর্বাভাস দেয়, তাই আপনি logits[:, :-1] (0 থেকে T-2 অবস্থানের পূর্বাভাস) এবং tokens[:, 1:] (1 থেকে T-1 অবস্থানের প্রকৃত টোকেন)-এর মধ্যে তুলনা করেন।

@nnx.jit
def train_step(model, optimizer, tokens):
    def loss_fn(model):
        logits = model(tokens)
        pred = logits[:, :-1].reshape(-1, VOCAB_SIZE)
        target = tokens[:, 1:].reshape(-1)
        return optax.softmax_cross_entropy_with_integer_labels(pred, target).mean()

    loss, grads = nnx.value_and_grad(loss_fn)(model)
    optimizer.update(model, grads)
    return {"loss": loss, "perplexity": jnp.exp(loss)}

ঐ ধাপের চারপাশের লুপটি একবার প্রস্তুতি নেয়, ফলে কম্পাইলেশনের সময় গণনা করা হয় না, তারপর এটি একটি নির্দিষ্ট সংখ্যক ধাপ চালায় এবং অতিবাহিত সময়কে টোকেন/সেকেন্ডে রূপান্তর করে।

def train_loop(model, optimizer, batches, steps=1000, log_every=100):
    num_batches = batches.shape[0]
    history = []
    # Warmup: compile the training step
    warmup_metrics = train_step(model, optimizer, batches[0])
    block_tree(warmup_metrics)
    start = time.perf_counter()
    for step in range(steps):
        tokens = batches[step % num_batches]
        metrics = train_step(model, optimizer, tokens)
        if step % log_every == 0 or step == steps - 1:
            metrics = block_tree(metrics)
            history.append(
                {
                    "step": step,
                    "loss": float(metrics["loss"]),
                    "perplexity": float(metrics["perplexity"]),
                }
            )
    block_tree(metrics)
    elapsed = time.perf_counter() - start
    batch_size = int(batches.shape[1])
    tokens_per_step = batch_size * (SEQ_LEN - 1)
    tokens_per_sec = steps * tokens_per_step / elapsed
    return history, elapsed, tokens_per_sec

এই সেলটিতে কেবল দুটি ফাংশন সংজ্ঞায়িত করা আছে, তাই এটি কোনো আউটপুট তৈরি করে না। পরবর্তী ধাপে এগুলো রান হয়।

৬. একটিমাত্র জিপিইউ-তে প্রশিক্ষণ দিন

একটি বেসলাইন স্থাপন করতে একটি GPU দিয়ে শুরু করুন। jax.device_put পুরো ব্যাচ অ্যারেটিকে gpu_devices[0] এ পিন করে, তাই ইচ্ছাকৃত মাল্টি-GPU তুলনার আগে কোনও কিছুই ডিভাইস জুড়ে ছড়িয়ে পড়ে না।

BENCHMARK_STEPS = 500
STEPS_1GPU = BENCHMARK_STEPS
batches_1gpu = make_batches(train_data, PER_DEVICE_BATCH)
single_device = gpu_devices[0]
batches_1gpu = jax.device_put(batches_1gpu, single_device)
model_1gpu = TinyTransformer(
    VOCAB_SIZE,
    D_MODEL,
    NUM_HEADS,
    FFN_DIM,
    NUM_LAYERS,
    MAX_SEQ_LEN,
    rngs=nnx.Rngs(1),
)
optimizer_1gpu = nnx.Optimizer(
    model_1gpu, optax.adamw(LR, weight_decay=WEIGHT_DECAY), wrt=nnx.Param
)
history_1gpu, elapsed_1gpu, tps_1gpu = train_loop(
    model_1gpu, optimizer_1gpu, batches_1gpu, steps=STEPS_1GPU
)
show_table(
    ["Step", "Loss", "Perplexity"],
    [(h["step"], f"{h['loss']:.3f}", f"{h['perplexity']:.1f}") for h in history_1gpu],
    title=f"Single-GPU training \u2014 {tps_1gpu:,.0f} tokens/sec",
    aligns=["right", "right", "right"],
)

সর্বপ্রথম কলটি স্টেপটি কম্পাইল করে, যে কারণে train_loop ঘড়ি চালু করার আগে ওয়ার্ম-আপ করে। রানটি শেষ হলে আপনি একটি টেবিল দেখতে পাবেন, যেখানে লগ করা প্রতিটি স্টেপের জন্য একটি করে সারি থাকবে এবং এতে দেখা যাবে যে ট্রেনিং এগোনোর সাথে সাথে লস ও পারপ্লেক্সিটি কমছে। টেবিলের শিরোনামে এই রানের জন্য পরিমাপ করা `tokens/sec` রিপোর্ট করা থাকে।

৭. সমস্ত জিপিইউ-তে একই ধাপ প্রয়োগ করুন

এটি একটি ডেটা-প্যারালাল প্যাটার্ন যেখানে আপনি একটি মেশ তৈরি করেন, মডেলটি রেপ্লিকেট করেন এবং ব্যাচ ডাইমেনশন বরাবর ডেটা শার্ড করেন। ট্রেনিং স্টেপের কোডে কোনো পরিবর্তন হয় না, কারণ অ্যারেগুলো কীভাবে সাজানো আছে তার উপর ভিত্তি করে @nnx.jit প্যারালালিজম পরিচালনা করে।

থ্রুপুট তুলনা ন্যায্য রাখার জন্য, সিঙ্গেল-জিপিইউ এবং মাল্টি-জিপিইউ রান একই সংখ্যক টাইমড স্টেপ ব্যবহার করে। প্রতিটি জিপিইউ প্রতি স্টেপে PER_DEVICE_BATCH সিকোয়েন্স প্রসেস করে, তাই মাল্টি-জিপিইউ রান একটি বৃহত্তর গ্লোবাল ব্যাচ প্রসেস করে।

একটি NNX মডেল রেপ্লিকেট করতে তিনটি কলের প্রয়োজন হয়। nnx.state মডিউলের স্টেটকে একটি PyTree হিসেবে এক্সট্র্যাক্ট করে, jax.device_put সেই PyTree-টিকে রেপ্লিকেটেড শার্ডিং সহ প্রতিটি ডিভাইসে স্থাপন করে, এবং nnx.update সেটিকে আবার মডিউলে লিখে দেয়। অপটিমাইজার স্টেটের ক্ষেত্রেও একই প্রক্রিয়া অনুসরণ করা হয়।

STEPS_MULTI = BENCHMARK_STEPS
GLOBAL_BATCH = PER_DEVICE_BATCH * NUM_DEVICES
mesh = Mesh(np.array(gpu_devices), ("data",))
replicated = NamedSharding(mesh, P())
data_sharding = NamedSharding(mesh, P(None, "data", None))
batches_multi = make_batches(train_data, GLOBAL_BATCH)
batches_multi = jax.device_put(batches_multi, data_sharding)
model_multi = TinyTransformer(
    VOCAB_SIZE,
    D_MODEL,
    NUM_HEADS,
    FFN_DIM,
    NUM_LAYERS,
    MAX_SEQ_LEN,
    rngs=nnx.Rngs(1),
)
optimizer_multi = nnx.Optimizer(
    model_multi, optax.adamw(LR, weight_decay=WEIGHT_DECAY), wrt=nnx.Param
)
# Replicate model and optimizer state across all GPUs
model_state = nnx.state(model_multi)
nnx.update(model_multi, jax.device_put(model_state, replicated))
opt_state = nnx.state(optimizer_multi)
nnx.update(optimizer_multi, jax.device_put(opt_state, replicated))
print(
    f"Global batch: {GLOBAL_BATCH} ({PER_DEVICE_BATCH} per GPU \u00d7 {NUM_DEVICES} GPUs)"
)
print(f"Training batches: {batches_multi.shape}")
print()
history_multi, elapsed_multi, tps_multi = train_loop(
    model_multi, optimizer_multi, batches_multi, steps=STEPS_MULTI
)
show_table(
    ["Step", "Loss", "Perplexity"],
    [(h["step"], f"{h['loss']:.3f}", f"{h['perplexity']:.1f}") for h in history_multi],
    title=f"Multi-GPU training \u2014 {tps_multi:,.0f} tokens/sec",
    aligns=["right", "right", "right"],
)

আপনি গ্লোবাল ব্যাচ লাইনটি দেখতে পাবেন, যা প্রতি GPU-এর জন্য PER_DEVICE_BATCH সিকোয়েন্সের সংখ্যাকে GPU-এর সংখ্যা দিয়ে গুণ করে, শার্ডেড ব্যাচ অ্যারের গঠন এবং তারপরে সিঙ্গেল-GPU রানের মতো একই কলামসহ একটি দ্বিতীয় লস টেবিল রিপোর্ট করে। এর শিরোনামে মাল্টি-GPU টোকেন/সেকেন্ডের সংখ্যা রিপোর্ট করা থাকে।

৮. থ্রুপুট তুলনা করুন এবং কার্ভগুলো অঙ্কন করুন।

উভয় রানেই মডেল এবং ট্রেনিং স্টেপ একই ছিল। শুধুমাত্র ডেটার বিন্যাস পরিবর্তিত হয়েছিল। উভয় রানেই একই সংখ্যক টাইমড স্টেপ ব্যবহৃত হয়েছিল, এবং প্রতিটি জিপিইউ প্রতি স্টেপে PER_DEVICE_BATCH সিকোয়েন্স প্রসেস করেছিল, তাই মাল্টি-জিপিইউ রানটির গ্লোবাল ব্যাচ বড় ছিল।

ms_per_step_1gpu = elapsed_1gpu / STEPS_1GPU * 1e3
ms_per_step_multi = elapsed_multi / STEPS_MULTI * 1e3
speedup = tps_multi / tps_1gpu

show_table(
    ["", "1 GPU", f"{NUM_DEVICES} GPUs", "Ratio"],
    [
        ("Batch size", PER_DEVICE_BATCH, GLOBAL_BATCH, f"{NUM_DEVICES}×"),
        ("Per-GPU batch", PER_DEVICE_BATCH, PER_DEVICE_BATCH, "same"),
        ("Timed steps", STEPS_1GPU, STEPS_MULTI, "same"),
        ("ms/step", f"{ms_per_step_1gpu:.2f}", f"{ms_per_step_multi:.2f}", f"{ms_per_step_1gpu / ms_per_step_multi:.2f}×"),
        ("Tokens/sec", f"{tps_1gpu:,.0f}", f"{tps_multi:,.0f}", f"{speedup:.2f}×"),
    ],
    title="Throughput comparison",
    aligns=["left", "right", "right", "right"],
)

if speedup > NUM_DEVICES * 1.25:
    print(
        f"Note: the measured speedup is superlinear (> {NUM_DEVICES}x). "
        "For this small benchmark, treat that as a measurement artifact rather "
        "than a general hardware-scaling claim."
    )

show_bars(
    [("1 GPU", tps_1gpu), (f"{NUM_DEVICES} GPUs", tps_multi)],
    "Training throughput",
    "tokens/s",
)

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

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

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
for label, history, batch_size in [
    ("1 GPU", history_1gpu, PER_DEVICE_BATCH),
    (f"{NUM_DEVICES} GPUs", history_multi, GLOBAL_BATCH),
]:
    tokens_m = [
        (h["step"] + 1) * batch_size * (SEQ_LEN - 1) / 1e6 for h in history
    ]
    losses = [h["loss"] for h in history]
    perps = [h["perplexity"] for h in history]
    ax1.plot(tokens_m, losses, "o-", label=label, markersize=4)
    ax2.plot(tokens_m, perps, "o-", label=label, markersize=4)
ax1.set_xlabel("Tokens processed (millions)")
ax1.set_ylabel("Loss")
ax1.set_title("Training loss")
ax1.legend()
ax1.grid(True, alpha=0.25)
ax2.set_xlabel("Tokens processed (millions)")
ax2.set_ylabel("Perplexity")
ax2.set_title("Training perplexity")
ax2.legend()
ax2.grid(True, alpha=0.25)
fig.suptitle("Training curves vs tokens processed")
fig.tight_layout()
plt.show()

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

৯. Orbax ব্যবহার করে একটি চেকপয়েন্ট সংরক্ষণ ও পুনরুদ্ধার করুন।

Orbax মডেলের অবস্থা অ্যারে ফাইলের একটি ডিরেক্টরি হিসেবে সংরক্ষণ করে। StandardCheckpointer হলো এর সবচেয়ে সরল API: সংরক্ষণের জন্য একটি কল, পুনরুদ্ধারের জন্য একটি কল।

NNX মডেলের ক্ষেত্রে, আপনি nnx.state(model, nnx.Param) ব্যবহার করে প্যারামিটারগুলো বের করেন, সেই PyTree-টি সেভ করেন এবং পরে nnx.update ব্যবহার করে একটি নতুন মডেলে তা পুনরুদ্ধার করেন। গ্রাফের কাঠামো ( nnx.GraphDef ) সেভ করা থাকে না এবং এটি পাইথন ক্লাস ডেফিনিশন থেকে আসে, তাই মডেলে ওয়েট লোড করার আগে মডেলটি পুনর্নির্মাণের জন্য আপনার TinyTransformer স্কোপের মধ্যে রাখতে হবে।

নিচের কোডটি সম্পূর্ণ লাইফসাইকেলটি চালায়। এটি শুধুমাত্র প্রশিক্ষিত মডেলের প্যারামিটারগুলো বের করে, সেগুলোকে ডিস্কে সংরক্ষণ করে, একটি ShapeDtypeStruct ট্রি তৈরি করে যা Orbax-কে জানায় কী ধরনের শেপ ও ডেটাটাইপ আশা করা যায়, একটি নতুনভাবে ইনিশিয়ালাইজ করা TinyTransformer এ ডেটা রিস্টোর করে, এবং একটি ছোট টেস্ট ইনপুটে পরীক্ষা করে দেখে যে রিস্টোর করা মডেলটি মূল মডেলের মতোই একই লগিট তৈরি করছে কিনা।

ckpt_dir = pathlib.Path("/tmp/jax-course/l7-checkpoints")
# Extract model parameters (not optimizer state)
model_params = nnx.state(model_multi, nnx.Param)
# Save
checkpointer = ocp.StandardCheckpointer()
if (ckpt_dir / "trained").exists():
    import shutil

    shutil.rmtree(ckpt_dir / "trained")
checkpointer.save(ckpt_dir / "trained", model_params)
print(f"Checkpoint saved to {ckpt_dir / 'trained'}")
# Create abstract target for restore
abstract_params = jax.tree.map(
    lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype),
    model_params,
)
# Restore into a fresh model
model_restored = TinyTransformer(
    VOCAB_SIZE,
    D_MODEL,
    NUM_HEADS,
    FFN_DIM,
    NUM_LAYERS,
    MAX_SEQ_LEN,
    rngs=nnx.Rngs(99),
)
restored_params = checkpointer.restore(ckpt_dir / "trained", abstract_params)
nnx.update(model_restored, restored_params)
# Test
test_input = jnp.zeros((1, 16), dtype=jnp.int32)
logits_original = model_multi(test_input)
logits_restored = model_restored(test_input)
max_diff = float(jnp.max(jnp.abs(logits_original - logits_restored)))
show_table(
    ["", "Value"],
    [
        ("Checkpoint path", str(ckpt_dir / "trained")),
        ("Parameters saved", f"{sum(x.size for x in jax.tree.leaves(model_params)):,}"),
        ("Max |original \u2212 restored|", f"{max_diff:.2e}"),
        ("Match", "\u2713" if max_diff < 1e-5 else "\u2717"),
    ],
    title="Orbax checkpoint save and restore",
)

আপনি চেকপয়েন্ট পাথ, সংরক্ষিত প্যারামিটারের সংখ্যা, মূল এবং পুনরুদ্ধার করা লজিটের মধ্যে সর্বোচ্চ পরম পার্থক্য এবং সেই পার্থক্য 1e-5 এর নিচে হলে একটি টিক চিহ্ন দেখতে পাবেন।

১০. শেক্সপিয়রের মতো লেখা তৈরি করুন

প্রশিক্ষিত মডেলটি প্রতিটি অবস্থানে পরবর্তী বাইটটি অনুমান করে। টেক্সট তৈরি করার জন্য, আপনি এটিকে একটি প্রম্পট দেন, শেষ অবস্থানের লগিটগুলো নেন, একটি টোকেন স্যাম্পল করেন, সেটিকে যুক্ত করেন এবং এই প্রক্রিয়াটি পুনরাবৃত্তি করেন।

generate ফাংশন তার ইনপুটকে MAX_SEQ_LEN পর্যন্ত প্যাড করে, ফলে সিকোয়েন্স বড় হওয়ার সাথে সাথে JIT-কম্পাইল করা ফরোয়ার্ড পাসে কোনো রিকম্পাইলেশন ছাড়াই সর্বদা একই ইনপুট শেপ দেখা যায়। কজাল অ্যাটেনশনের ক্ষেত্রে, আসল টোকেনগুলোর পরে প্যাডিং করা হলে তা আগের পজিশনগুলোর আউটপুটকে প্রভাবিত করে না।

nnx.split মডিউলটিকে একটি graphdef এবং একটি state PyTree-তে বিভক্ত করে, যাতে এটিকে jax.jit মাধ্যমে পাস করা যায়, এবং nnx.merge কম্পাইল করা ফাংশনের ভিতরে মডিউলটিকে পুনর্নির্মাণ করে।

@jax.jit
def get_logits_jit(graphdef, model_state, tokens):
    model = nnx.merge(graphdef, model_state)
    return model(tokens)

def generate(model, prompt_text, max_new_tokens=300, temperature=0.8):
    graphdef, model_state = nnx.split(model)
    tokens = list(prompt_text.encode("utf-8"))
    key = jax.random.key(42)
    for _ in range(max_new_tokens):
        context = tokens[-MAX_SEQ_LEN:]
        padded = context + [0] * (MAX_SEQ_LEN - len(context))
        input_arr = jnp.array([padded], dtype=jnp.int32)
        logits = get_logits_jit(graphdef, model_state, input_arr)
        next_logit = logits[0, len(context) - 1]
        if temperature <= 0:
            next_token = int(jnp.argmax(next_logit))
        else:
            key, subkey = jax.random.split(key)
            next_token = int(jax.random.categorical(subkey, next_logit / temperature))
        tokens.append(next_token)
    return bytes(tokens).decode("utf-8", errors="replace")

# Put model on a single device for generation
gen_model = TinyTransformer(
    VOCAB_SIZE,
    D_MODEL,
    NUM_HEADS,
    FFN_DIM,
    NUM_LAYERS,
    MAX_SEQ_LEN,
    rngs=nnx.Rngs(99),
)
nnx.update(gen_model, checkpointer.restore(ckpt_dir / "trained", abstract_params))
print("=== Prompt: 'ROMEO:' | temperature=0.8 ===")
print()
print(generate(gen_model, "ROMEO:", max_new_tokens=300, temperature=0.8))
print()
print("=== Prompt: 'To be, or not' | temperature=0.6 ===")
print()
print(generate(gen_model, "To be, or not", max_new_tokens=300, temperature=0.6))

আপনি তৈরি হওয়া টেক্সটের দুটি ব্লক দেখতে পাবেন, প্রতিটি প্রম্পটের জন্য একটি করে।

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

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

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

kubectl delete -f deploy/jupyter.yaml

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

cd terraform
terraform destroy

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

gcloud container clusters list
gcloud compute instances list

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

১২. অভিনন্দন

আপনি Flax NNX ব্যবহার করে একটি ডিকোডার ট্রান্সফরমার তৈরি করেছেন, আপনার নোডের উভয় GPU-তে এটিকে প্রশিক্ষণ দিয়েছেন, Orbax দিয়ে এর চেকপয়েন্ট নির্ধারণ করেছেন এবং পুনরুদ্ধার করা ওয়েটগুলো থেকে টেক্সট তৈরি করেছেন।

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

• Flax NNX যেভাবে একটি মডেলকে পুনঃব্যবহারযোগ্য মডিউলে সংগঠিত করে: nnx.Embed , nnx.Linear , nnx.LayerNorm , এবং nnx.MultiHeadAttention প্যারামিটার তৈরি এবং ফরোয়ার্ড পাস পরিচালনা করে, এবং ট্রেনিং ধাপে nnx.Optimizer এর সাথে nnx.value_and_grad ব্যবহৃত হয়।
• কিভাবে is_causal=True সহ কজাল অ্যাটেনশন ভবিষ্যতের অবস্থানগুলিকে মাস্ক করে যাতে মডেলটি কেবল পিছনে তাকাতে পারে, এবং কেন attention_fn হুকের একটি **_ র‍্যাপার প্রয়োজন
• কোডল্যাব ৬-এর মতোই, ট্রেনিং স্টেপ পরিবর্তন না করেই ডেটা-প্যারালাল ট্রেনিং কীভাবে মডেলটিকে প্রতিলিপি করে এবং জিপিইউ জুড়ে ব্যাচকে বিভক্ত করে।
• nnx.state এবং nnx.update কীভাবে NNX মডিউল এবং সাধারণ PyTree-এর মধ্যে সংযোগ স্থাপন করে, তার মাধ্যমে Orbax কীভাবে মডেল প্যারামিটার সংরক্ষণ ও পুনরুদ্ধার করে।
• ল্যাঙ্গুয়েজ মডেলের স্বাভাবিক একক টোকেন/সেকেন্ড-এ থ্রুপুট কীভাবে পরিমাপ করা যায়, যেখানে ওয়ার্মআপ এবং block_until_ready সংখ্যাগুলোকে নির্ভুল রাখার কাজটি চালিয়ে যায়।
• কীভাবে টেক্সট জেনারেশন মডেলকে একবারে একটি টোকেন করে তার নিজস্ব পূর্বাভাস সরবরাহ করে, এবং কীভাবে তাপমাত্রা সেই স্যাম্পলিংয়ের এলোমেলো ভাবকে নিয়ন্ত্রণ করে।

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

• কোডল্যাব ৮: একটি প্রশিক্ষিত JAX মডেল এক্সপোর্ট এবং সার্ভ করুন, যেখানে আপনি এই চেকপয়েন্টটি ব্যবহার করে JIT ইনফারেন্স, AOT কম্পাইলেশন এবং পোর্টেবল ফরম্যাটে এক্সপোর্টের মাধ্যমে এটিকে সার্ভ করার জন্য প্রস্তুত করবেন।
• D_MODEL (১২৮ বা ৫১২ চেষ্টা করুন) এবং NUM_LAYERS (২ বা ৬ চেষ্টা করুন) পরিবর্তন করে আবার চালান। ধারণক্ষমতা বেশি হলে ধাপগুলো ধীর হবে, এবং লেয়ার বেশি হলে বেশি মেমরি লাগবে।
• জেনারেশন টেম্পারেচার পরিবর্তন করুন (০.০, ০.৫, এবং ১.২ চেষ্টা করে দেখুন)। শূন্য হলো গ্রিডি, এবং ১-এর উপরের মানগুলো আরও র‍্যান্ডম আউটপুট তৈরি করে।
• অ্যাটেনশন ব্যাকএন্ড পরিবর্তন করুন: causal_sdpa এর বডি পরিবর্তন করে jax.nn.dot_product_attention(..., is_causal=True, implementation="cudnn") কল করুন এবং bf16-সামঞ্জস্যপূর্ণ অ্যাক্টিভেশন ব্যবহার করুন, তারপর দেখুন cuDNN ফিউজড অ্যাটেনশন কীভাবে আচরণ করে।

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

• ফ্ল্যাক্স এনএনএক্স এর প্রাথমিক বিষয়াবলী
• Flax NNX এবং JAX রূপান্তর
• অরবাক্স চেকপয়েন্ট গাইড
• nnx.MultiHeadAttention API
• অপট্যাক্স অ্যাডামডব্লিউ
• টাইনিশেক্সপিয়ার ডেটাসেট
• NGC-তে NVIDIA JAX কন্টেইনার