একটি প্রশিক্ষিত JAX মডেল রপ্তানি এবং পরিবেশন করুন

১. ভূমিকা

এই কোডল্যাবে আপনি একটি প্রশিক্ষিত চেকপয়েন্ট থেকে শুরু করে JAX ইনফারেন্স পাইপলাইনের শুরু থেকে শেষ পর্যন্ত অনুসরণ করবেন; একটি JIT-কম্পাইলড ফরোয়ার্ড পাস থেকে শুরু করে অ্যাহেড-অফ-টাইম (AOT) কম্পাইলেশন, jax.export ব্যবহার করে JAX-নেটিভ এক্সপোর্ট এবং jax2tf ব্যবহার করে একটি TensorFlow SavedModel পর্যন্ত। প্রতিটি পথ পরিমাপ করা হয় এবং এই চারটি পথের উদ্দেশ্য হলো একই প্রেডিকশন তৈরি করা।

আপনি যা করবেন

• কোডল্যাব ৭ ট্রান্সফর্মারটি পুনর্নির্মাণ করুন এবং Orbax ব্যবহার করে এর প্রশিক্ষিত ওয়েটগুলো লোড করুন।
• ফরওয়ার্ড পাসটিকে jax.jit মধ্যে আবদ্ধ করুন এবং ক্যাশড-কল ল্যাটেন্সির সাথে ফার্স্ট-কল ল্যাটেন্সির পরিমাপ করুন।
• AOT কম্পাইলেশন ( lower() তারপর compile() ) ব্যবহার করে কোল্ড স্টার্ট দূর করুন এবং StableHLO IR পড়ুন।
• চারটি ব্যাচ সাইজ জুড়ে টোকেন/সেকেন্ড এককে ফরোয়ার্ড-পাস থ্রুপুট পরিমাপ করুন।
• jax.export ব্যবহার করে একটি পোর্টেবল আর্টিফ্যাক্ট সিরিয়ালাইজ করুন, তারপর ডিসিরিয়ালাইজ করে সেটিকে কল করুন।
• jax2tf ব্যবহার করে মডেলটিকে একটি TensorFlow SavedModel-এ রূপান্তর করুন এবং চারটি পাথই তুলনা করুন।

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

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

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

প্রশিক্ষণ থেকে সেবায় যাওয়ার চারটি পথ

এই কোডল্যাবটি আপনার ডেপ্লয়মেন্ট টার্গেটের উপর নির্ভর করে, একটি প্রশিক্ষিত JAX মডেলকে সার্ভিংয়ের দিকে নিয়ে যাওয়ার চারটি কার্যকরী পদ্ধতির তুলনা করে:

পথ	ফর্ম্যাট	পরিবেশন লক্ষ্য	কখন ব্যবহার করবেন
jax.jit	ক্যাশে করা চলমান নির্বাহক	পাইথন সার্ভার (ফাস্টএপিআই, ফ্লাস্ক)	সবচেয়ে সহজ কম-লেটেন্সি JAX সার্ভিং পাথ
AOT compile	প্রি-কম্পাইলড ইন-প্রসেস এক্সিকিউটেবল	পাইথন সার্ভার চালু / ওয়ার্মআপ	প্রথম অনুরোধের কম্পাইলেশন লেটেন্সি এড়িয়ে চলুন
jax.export	StableHLO + মেটাডেটা সহ সিরিয়ালাইজড JAX এক্সপোর্ট	রপ্তানিকৃত প্ল্যাটফর্ম(গুলি)র জন্য সামঞ্জস্যপূর্ণ JAX রানটাইম	JAX-নেটিভ পোর্টেবল আর্টিফ্যাক্ট
jax2tf	TF সংরক্ষিত মডেল	টিএফ সার্ভিং, টিএফএক্স	TensorFlow ইকোসিস্টেম

আপনি একই প্রশিক্ষিত চেকপয়েন্ট থেকে শুরু করে চারটিই চালান। পার্থক্যগুলো হলো বহনযোগ্যতা, নির্ভরতা এবং আর্টিফ্যাক্টটি মূল পাইথন প্রক্রিয়ার বাইরে ব্যবহার করা যাবে কিনা।

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

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

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

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

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

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 orbax-checkpoint tensorflow

flax এবং orbax-checkpoint সাধারণত NVIDIA JAX কন্টেইনারের সাথেই আসে, তাই এই দুটি সাধারণত কোনো কাজ করে না। যদি TensorFlow আগে থেকে উপস্থিত না থাকে, pip এটিকে পুল করে নেয়।

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

JAX, Flax NNX, এবং Orbax ইম্পোর্ট করুন, নিশ্চিত করুন যে অন্তত একটি GPU দৃশ্যমান আছে, এবং ফলাফলের জন্য ব্লক করতে ও ফলাফল টেবিল রেন্ডার করতে কোডল্যাবের বাকি অংশে ব্যবহৃত দুটি হেল্পার সংজ্ঞায়িত করুন।

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

import html
import pathlib
import time
import warnings

from IPython.display import HTML, display
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
from flax import nnx
import orbax.checkpoint as ocp


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

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

assert len(gpu_devices) >= 1, (
    f"This lesson needs at least 1 GPU. 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("\n".join(parts)))

আপনি একটি JAX ভার্সন, ডিফল্ট ব্যাকএন্ড হিসেবে gpu , এবং অন্তত একটি CUDA ডিভাইসসহ একটি তালিকা দেখতে পাবেন।

আপনার একটি প্রশিক্ষিত চেকপয়েন্ট আছে কিনা তা নিশ্চিত করুন।

এই কোডল্যাবটি Orbax-এর সেই চেকপয়েন্টটি লোড করে যা 'Train a transformer end to end with Flax NNX and Orbax' কোডল্যাবটি /tmp/jax-course/l7-checkpoints/trained ডিরেক্টরিতে লিখেছিল, তাই কিছু বিল্ড করার আগে ডিরেক্টরিটি এখনও আছে কিনা তা দেখে নিন।

import pathlib
ckpt_dir = pathlib.Path("/tmp/jax-course/l7-checkpoints")
assert (ckpt_dir / "trained").exists(), (
    f"No checkpoint at {ckpt_dir / 'trained'}. Run codelab 7 first."
)
print(f"Found checkpoint: {ckpt_dir / 'trained'}")

আপনার কাছে চেকপয়েন্ট পাথটি প্রিন্ট হয়ে ফিরে আসবে।

৩. মডেলটি পুনর্নির্মাণ করুন এবং চেকপয়েন্ট লোড করুন।

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

স্থাপত্যকে নতুনভাবে সংজ্ঞায়িত করুন

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

VOCAB_SIZE = 256
D_MODEL = 256
NUM_HEADS = 4
FFN_DIM = 1024
NUM_LAYERS = 4
MAX_SEQ_LEN = 256


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


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


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)


param_count = sum(x.size for x in jax.tree.leaves(nnx.state(TinyTransformer(
    VOCAB_SIZE, D_MODEL, NUM_HEADS, FFN_DIM, NUM_LAYERS, MAX_SEQ_LEN,
    rngs=nnx.Rngs(0),
), nnx.Param)))
print(f"TinyTransformer: {param_count:,} parameters")

একটি সদ্য চালু করা মডেলের প্যারামিটার সংখ্যা আপনি দেখতে পাবেন। এই পর্যায়ে ওয়েটগুলো এলোমেলো থাকে, কারণ এখানে আর্কিটেকচারটাই গুরুত্বপূর্ণ।

প্রশিক্ষিত ওজন পুনরুদ্ধার করুন

Orbax একটি টার্গেট স্ট্রাকচারে রিস্টোর করে। নতুন মডেলের স্টেট নিয়ে এবং প্রতিটি অ্যারে-কে একটি jax.ShapeDtypeStruct দিয়ে প্রতিস্থাপন করে সেই টার্গেটটি তৈরি করা হয়, যা কোনো মেমরি অ্যালোকেট না করেই একটি শেপ এবং ডেটাটাইপ বর্ণনা করে। এরপর Orbax ঠিক সেই শেপগুলোর অ্যারেতে চেকপয়েন্টটি রিড করে নেয়।

ckpt_dir = pathlib.Path("/tmp/jax-course/l7-checkpoints")

model = TinyTransformer(
    VOCAB_SIZE, D_MODEL, NUM_HEADS, FFN_DIM, NUM_LAYERS, MAX_SEQ_LEN,
    rngs=nnx.Rngs(0),
)
model_params = nnx.state(model, nnx.Param)
abstract_params = jax.tree.map(
    lambda x: jax.ShapeDtypeStruct(x.shape, x.dtype), model_params
)

checkpointer = ocp.StandardCheckpointer()
restored_params = checkpointer.restore(ckpt_dir / "trained", abstract_params)

# Move to a single GPU
single_device = gpu_devices[0]
restored_params = jax.device_put(restored_params, single_device)
nnx.update(model, restored_params)

print(f"\u2705 Checkpoint loaded from {ckpt_dir / 'trained'}")
print(f"Parameters: {sum(x.size for x in jax.tree.leaves(restored_params)):,}")

jax.device_put কলটি গুরুত্বপূর্ণ। সার্ভিং একটি GPU ব্যবহার করে, তাই প্যারামিটারগুলি মডেলে যাওয়ার আগে gpu_devices[0] এ সরানো হয়।

আপনি চেকপয়েন্ট পাথটি নিশ্চিত হতে দেখবেন এবং প্যারামিটার সংখ্যাটি পূর্ববর্তী সেল দ্বারা প্রিন্ট করা সংখ্যার সাথে মিলে যাবে।

৪. পথ ১: jax.jit দিয়ে পরিবেশন করুন

সবচেয়ে সহজ সার্ভিং পদ্ধতি হলো ফরওয়ার্ড পাসটিকে jax.jit এর মধ্যে র‍্যাপ করা। প্রথম কলটি কম্পাইলেশন শুরু করে এবং এর পরের প্রতিটি কল ক্যাশ করা এক্সিকিউটেবলটি পুনরায় ব্যবহার করে।

সার্ভ করার জন্য, ওয়েটগুলোকে একটি ক্লোজারের মধ্যে বেক করুন। nnx.split মডেলটিকে একটি স্ট্যাটিক গ্রাফ ডেফিনিশন এবং ডাইনামিক স্টেটে বিভক্ত করে, এবং jitted ফাংশনটি উভয়কেই ধারণ করে ও আর্গুমেন্ট হিসেবে শুধুমাত্র tokens গ্রহণ করে। এর ফলে কম্পাইল করা ফাংশনটি স্বয়ংসম্পূর্ণ হয়।

graphdef, model_state = nnx.split(model)

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

dummy_input = jnp.zeros((1, MAX_SEQ_LEN), dtype=jnp.int32)

start = time.perf_counter()
logits = block_tree(predict_jit(dummy_input))
first_call_ms = (time.perf_counter() - start) * 1000

times = []
for _ in range(100):
    start = time.perf_counter()
    logits = block_tree(predict_jit(dummy_input))
    times.append((time.perf_counter() - start) * 1000)

avg_ms = np.mean(times)

show_table(
    ["", "Latency (ms)"],
    [
        ("First call (compile + execute)", f"{first_call_ms:,.1f}"),
        ("Subsequent calls (avg of 100)", f"{avg_ms:.2f}"),
        ("Speedup", f"{first_call_ms / avg_ms:.0f}\u00d7"),
    ],
    title="JIT inference latency",
    aligns=["left", "right"],
)

print(f"\nOutput shape: {logits.shape} (batch=1, seq={MAX_SEQ_LEN}, vocab={VOCAB_SIZE})")

আপনি একটি JIT ইনফারেন্স ল্যাটেন্সি টেবিল দেখতে পাবেন, যেখানে প্রথম কলটি পরবর্তী ১০০টির গড়ের চেয়ে অনেক ধীরগতির হবে এবং Output shape: (1, 256, 256) ।

৫. পথ ২: AOT কম্পাইলেশনের মাধ্যমে কোল্ড স্টার্ট দূর করুন

jax.jit প্রথম কলেই কম্পাইল হয়, যা ডেভেলপমেন্টের সময় ঠিক আছে কিন্তু সার্ভিংয়ের ক্ষেত্রে একটি সমস্যা। AOT কম্পাইলেশন সেই একক ধাপটিকে আলাদা আলাদা পর্যায়ে বিভক্ত করে:

1. আপনি একটি সাধারণ পাইথন/জ্যাক্স ফাংশন লেখেন, যেমন একটি মডেল ফরোয়ার্ড পাস
2. JAX নির্দিষ্ট ইনপুট আকার এবং ডেটাটাইপের জন্য ফাংশনটি ট্রেস করে এবং lower() ব্যবহার করে এটিকে কম্পাইলার IR-এ রূপান্তর করে।
3. StableHLO dot , reshape এবং reduce মতো হার্ডওয়্যার-নিরপেক্ষ অপারেশনগুলোর মাধ্যমে গণনা বর্ণনা করে।
4. XLA compile() মাধ্যমে StableHLO-কে অপ্টিমাইজ করে এবং GPU, TPU বা CPU-এর জন্য একটি ডিভাইস-নির্দিষ্ট এক্সিকিউটেবল তৈরি করে।

StableHLO ইন্টারমিডিয়েট রিপ্রেজেন্টেশন (IR) হলো একটি JAX কম্পিউটেশনের কম্পাইলার-স্তরের উপস্থাপনা, যা পাইথন কোডকে একটি পোর্টেবল ও হার্ডওয়্যার-নিরপেক্ষ প্রোগ্রামে রূপান্তর করার পর পাওয়া যায়। এটি ম্যাট্রিক্স গুণ, রিসেপ, রিডাকশন এবং কন্ট্রোল ফ্লো-এর মতো অপারেশনগুলোকে এমন একটি রূপে বর্ণনা করে, যা XLA বিভিন্ন ব্যাকএন্ড, যেমন GPU, TPU এবং CPU-এর জন্য কম্পাইল করতে পারে।

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

# Bake weights into a closure
def predict_closed(tokens):
    m = nnx.merge(graphdef, model_state)
    return m(tokens)

# Stage 1: Lower
abstract_tokens = jax.ShapeDtypeStruct((1, MAX_SEQ_LEN), jnp.int32)

lowered = predict_closed.lower(abstract_tokens)
print(f"Lowered to StableHLO ({len(lowered.as_text()):,} chars)")

# Stage 2: Compile
compiled = lowered.compile()
print(f"Compiled for: {jax.default_backend()}")

# Execute
start = time.perf_counter()
logits_aot = block_tree(compiled(dummy_input))
aot_first_ms = (time.perf_counter() - start) * 1000

times_aot = []
for _ in range(100):
    start = time.perf_counter()
    logits_aot = block_tree(compiled(dummy_input))
    times_aot.append((time.perf_counter() - start) * 1000)

avg_aot_ms = np.mean(times_aot)
max_diff_jit_aot = float(jnp.max(jnp.abs(logits - logits_aot)))

show_table(
    ["", "Latency (ms)"],
    [
        ("AOT first execution (no compile)", f"{aot_first_ms:.2f}"),
        ("AOT subsequent (avg of 100)", f"{avg_aot_ms:.2f}"),
        ("JIT first call (from above)", f"{first_call_ms:,.1f}"),
        ("Max |JIT − AOT|", f"{max_diff_jit_aot:.2e}"),
    ],
    title="AOT vs JIT latency",
    aligns=["left", "right"],
)

মনে রাখবেন যে lower() আসল ডেটা নয়, বরং একটি jax.ShapeDtypeStruct গ্রহণ করে। কম্পাইল করার জন্য আপনার ইনপুট অ্যারের কখনোই প্রয়োজন হয় না — শুধুমাত্র এর shape এবং dtype-ই যথেষ্ট।

আপনি এমন একটি AOT বনাম JIT লেটেন্সি টেবিল দেখতে পাবেন, যেখানে AOT-এর প্রথম এক্সিকিউশনটি JIT-এর প্রথম কলের পরিবর্তে AOT-এর স্টেডি-স্টেট সংখ্যার কাছাকাছি থাকে এবং Max |JIT − AOT| এর মান কার্যত শূন্য হয়।

StableHLO IR পরিদর্শন করুন

lowered.as_text() সেই StableHLO প্রোগ্রামটি দেখায় যা ডিভাইসে রান করবে। এটি সেই একই অন্তর্বর্তী উপস্থাপনা যা XLA, GPU, TPU এবং CPU জুড়ে ব্যবহার করে, এবং এটি ডিবাগিং, পারফরম্যান্স বিশ্লেষণ এবং কম্পাইলার আসলে কী দেখছে তা বোঝার জন্য উপযোগী।

StableHLO টেক্সট বেশ বড় হতে পারে, এবং একটি লাইনে একটি দীর্ঘ কনস্ট্যান্ট বা অ্যাট্রিবিউট থাকতে পারে। Jupyter-এর IOPub ডেটা-রেটের সীমাবদ্ধতা এড়ানোর জন্য, পরবর্তী কোডটি সম্পূর্ণ IR-টি ডিস্কে সংরক্ষণ করে এবং শুধুমাত্র একটি সীমাবদ্ধ প্রিভিউ প্রিন্ট করে।

hlo_text = lowered.as_text()

hlo_path = pathlib.Path("/tmp/jax-course/l8-stablehlo.mlir")
hlo_path.parent.mkdir(parents=True, exist_ok=True)
hlo_path.write_text(hlo_text)

MAX_LINES = 20
MAX_CHARS_PER_LINE = 160

lines = hlo_text.splitlines()
preview_lines = []
for line in lines[:MAX_LINES]:
    if len(line) > MAX_CHARS_PER_LINE:
        preview_lines.append(line[:MAX_CHARS_PER_LINE] + " ... [line truncated]")
    else:
        preview_lines.append(line)

print(f"StableHLO program: {len(lines):,} lines, {len(hlo_text):,} chars")
print(f"Full StableHLO saved to: {hlo_path}")
print("=" * 60)
print("\n".join(preview_lines))
print(
    f"\n... ({max(len(lines) - MAX_LINES, 0):,} more lines; "
    "long lines are truncated in this preview)"
)

আপনি প্রোগ্রামটির লাইন ও ক্যারেক্টার সংখ্যা, সংরক্ষিত .mlir ফাইলের পাথ এবং IR-এর প্রথম ২০টি লাইন দেখতে পাবেন, যার পরে বাদ দেওয়া লাইনের সংখ্যাও দেখানো হবে।

৬. ব্যাচড ইনফারেন্স থ্রুপুট পরিমাপ করুন

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

batch_sizes = [1, 4, 16, 64]
results = []

for bs in batch_sizes:
    tokens_batch = jnp.zeros((bs, MAX_SEQ_LEN), dtype=jnp.int32)

    # Compile for this batch size
    lowered_bs = predict_closed.lower(
        jax.ShapeDtypeStruct((bs, MAX_SEQ_LEN), jnp.int32),
    )
    compiled_bs = lowered_bs.compile()

    # Warmup
    block_tree(compiled_bs(tokens_batch))

    # Measure
    times_bs = []
    for _ in range(50):
        start = time.perf_counter()
        block_tree(compiled_bs(tokens_batch))
        times_bs.append((time.perf_counter() - start) * 1000)

    avg_bs = np.mean(times_bs)
    tokens_per_sec = (bs * MAX_SEQ_LEN) / (avg_bs / 1000)
    results.append((bs, f"{avg_bs:.2f}", f"{tokens_per_sec:,.0f}"))

show_table(
    ["Batch size", "Latency (ms)", "Tokens/sec"],
    results,
    title="Batched inference throughput",
    aligns=["right", "right", "right"],
)

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

৭. পদ্ধতি ৩: jax.export ব্যবহার করে একটি পোর্টেবল আর্টিফ্যাক্ট এক্সপোর্ট করুন।

jax.export একটি জিটেড JAX ফাংশনকে একটি Exported অবজেক্টে এক্সপোর্ট করে, যাতে StableHLO-এর পাশাপাশি অন্য JAX প্রসেস থেকে এটিকে কল করার জন্য প্রয়োজনীয় মেটাডেটা থাকে। সিরিয়ালাইজড বাইটগুলো হতে পারে:

• ডিস্কে সংরক্ষিত এবং অন্য একটি প্রক্রিয়ায় লোড করা হয়েছে
• মূল পাইথন মডেল সোর্স কোড ছাড়াই কল করা হয়েছে
• ডিফল্টরূপে বর্তমান প্ল্যাটফর্মের জন্য, অথবা platforms=[...] আর্গুমেন্টের মাধ্যমে সুনির্দিষ্ট প্ল্যাটফর্মের জন্য এক্সপোর্ট করা হয়।

এটি JAX-ভিত্তিক ডেপ্লয়মেন্ট পদ্ধতি — এর জন্য TensorFlow-এর উপর কোনো নির্ভরতার প্রয়োজন নেই।

from jax import export

# Export the closure-based function
exported = export.export(predict_closed)(
    jax.ShapeDtypeStruct((1, MAX_SEQ_LEN), jnp.int32),
)

print(f"Exported function: {exported.fun_name}")
print(f"Input shapes:  {exported.in_avals}")
print(f"Output shapes: {exported.out_avals}")
print(f"Exported platforms: {exported.platforms}")

# Serialize to bytes
blob = exported.serialize()
export_path = pathlib.Path("/tmp/jax-course/exports")
export_path.mkdir(parents=True, exist_ok=True)

export_file = export_path / "tiny_transformer_jax_export.bin"
export_file.write_bytes(blob)
print()
print(f"Serialized to {export_file} ({len(blob):,} bytes, {len(blob) / 1024:.0f} KB)")

# Deserialize and call
rehydrated = export.deserialize(export_file.read_bytes())

test_input = jnp.zeros((1, MAX_SEQ_LEN), dtype=jnp.int32)
logits_exported = block_tree(rehydrated.call(test_input))
print()
print(f"✅ Deserialized call succeeded — output shape: {logits_exported.shape}")

# Verify outputs match
diff_export = float(jnp.max(jnp.abs(logits_aot - logits_exported)))
print(f"Max difference from AOT: {diff_export:.2e}")

rehydrated অবজেক্টটি কখনোই TinyTransformer সংস্পর্শে আসেনি। এটিকে ডিস্কের বাইট থেকে পুনর্নির্মাণ করা হয়েছিল এবং এটি এখনও একই লগিট ফেরত দেয়, যা একটি পোর্টেবল আর্টিফ্যাক্টের মূল উদ্দেশ্য।

আপনি এক্সপোর্ট করা ফাংশনের নাম, ইনপুট ও আউটপুট অ্যাভাল, এক্সপোর্ট করা প্ল্যাটফর্ম তালিকা, সিরিয়ালাইজড ব্লবের বাইট সংখ্যা, ডিসিরিয়ালাইজড কল থেকে প্রাপ্ত আউটপুট শেপ এবং AOT লজিটস থেকে প্রাপ্ত সর্বোচ্চ পার্থক্য দেখতে পাবেন, যা কার্যত শূন্য।

৮. পথ ৪: jax2tf ব্যবহার করে TensorFlow SavedModel-এ রূপান্তর করুন

আপনার সার্ভিং পরিকাঠামো যদি TF Serving বা TFX পাইপলাইনের মতো TensorFlow ব্যবহার করে, তাহলে আপনি JAX ফাংশনটিকে একটি TF SavedModel-এ রূপান্তর করতে পারেন। jax2tf এখনও jax.experimental অধীনেই থাকে, কিন্তু এটিই JAX থেকে TensorFlow-এর মধ্যে যোগাযোগের আদর্শ পথ।

বর্তমান JAX রিলিজগুলিতে নেটিভ সিরিয়ালাইজেশনই ডিফল্ট, তাই native_serialization=True পাস না করেই এই রূপান্তরটি TensorFlow গ্রাফে লোয়ার্ড StableHLO অন্তর্ভুক্ত করে।

এখানে একটি প্ল্যাটফর্মের বিবরণ গুরুত্বপূর্ণ। এই কোডল্যাবটি CUDA-তে JAX চালায়, কিন্তু JAX কন্টেইনারের TensorFlow রানটাইম SavedModel-টিকে CPU-তে এক্সিকিউট করতে পারে। CPU-তে TensorFlow দ্বারা একটি CUDA-এক্সপোর্টেড মডিউল কল হওয়া এড়ানোর জন্য, নিচের কোডটি jax2tf মডিউলটিকে ("cpu",) এর জন্য এক্সপোর্ট করে। আপনি যদি GPU-তে TensorFlow দিয়ে সার্ভ করেন, তাহলে ("cuda",) এর জন্য এক্সপোর্ট করুন এবং GPU/XLA সাপোর্টসহ একটি TensorFlow রানটাইম ব্যবহার করুন।

from jax.experimental import jax2tf
import tensorflow as tf
import shutil

# Capture model_state as a closure
def predict_for_tf(tokens):
    m = nnx.merge(graphdef, model_state)
    return m(tokens)

TF_EXPORT_PLATFORMS = ("cpu",)
tf_predict = jax2tf.convert(
    predict_for_tf,
    native_serialization_platforms=TF_EXPORT_PLATFORMS,
)

# Wrap in a tf.Module for SavedModel export
module = tf.Module()
module.predict = tf.function(
    tf_predict,
    input_signature=[tf.TensorSpec(shape=(1, MAX_SEQ_LEN), dtype=tf.int32)],
    autograph=False,
)

# TF Serving expects a versioned model directory
savedmodel_base_dir = pathlib.Path("/tmp/jax-course/exports") / "tiny_transformer_savedmodel"
savedmodel_dir = savedmodel_base_dir / "1"
if savedmodel_base_dir.exists():
    shutil.rmtree(savedmodel_base_dir)

tf.saved_model.save(module, str(savedmodel_dir))
print(f"✅ SavedModel saved to {savedmodel_dir}")
print(f"Exported for TensorFlow platform(s): {TF_EXPORT_PLATFORMS}")

# Verify the SavedModel path produces the same logits as the JAX path
with tf.device("/CPU:0"):
    tf_logits = module.predict(tf.zeros((1, MAX_SEQ_LEN), dtype=tf.int32))
diff_tf = np.max(np.abs(np.asarray(tf_logits) - np.asarray(logits_aot)))
print(f"Max difference from AOT: {diff_tf:.2e}")

# List saved files
for dirpath, _, filenames in os.walk(savedmodel_base_dir):
    for f in filenames:
        full = os.path.join(dirpath, f)
        size = os.path.getsize(full)
        print(f"  {os.path.relpath(full, savedmodel_base_dir):44s} {size:>10,} bytes")
print()
print(
    f"To serve: docker run -p 8501:8501 "
    f"--mount type=bind,source={savedmodel_base_dir},target=/models/transformer "
    "-e MODEL_NAME=transformer tensorflow/serving "
    "--xla_cpu_compilation_enabled=true"
)

আপনি SavedModel ডিরেক্টরির পাথ, এক্সপোর্ট প্ল্যাটফর্ম টাপল, AOT লজিটস থেকে কার্যত শূন্য একটি সর্বোচ্চ পার্থক্য এবং ভার্সনযুক্ত ডিরেক্টরির অধীনে লেখা ফাইলগুলির একটি তালিকা দেখতে পাবেন।

কোডটি সবশেষে একটি docker run ... tensorflow/serving কমান্ড প্রিন্ট করে। এটি একটি দৃষ্টান্তমূলক বিষয় এবং এটি দেখায় যে, আপনি কীভাবে আপনার সদ্য তৈরি করা ভার্সনযুক্ত ডিরেক্টরিটি TensorFlow Serving-এর দিকে নির্দেশ করবেন। এই কোডল্যাবে এটি রান করা হয় না।

৯. চারটি পরিবেশন পথের তুলনা করুন

প্রতিটি পাথ একই প্রশিক্ষিত চেকপয়েন্ট থেকে শুরু হয়েছিল এবং একই পূর্বাভাস তৈরি করেছিল। পার্থক্য হলো বহনযোগ্যতা, নির্ভরতা, ডেপ্লয়মেন্টের লক্ষ্য এবং আর্টিফ্যাক্টটি মূল পাইথন প্রসেসের বাইরে ব্যবহার করা যাবে কিনা।

show_table(
    ["Path", "Format", "Dependencies", "Serving target", "Portable"],
    [
        ("jax.jit", "Cached in-process executable", "JAX", "Python server", "No"),
        ("AOT compile", "Compiled in-process executable", "JAX", "Python server startup / warmup", "No"),
        ("jax.export", "Serialized JAX export", "JAX runtime", "Compatible runtime for exported platform(s)", "Yes"),
        ("jax2tf", "TF SavedModel", "TensorFlow", "TF Serving", "Yes"),
    ],
    title="Serving path comparison",
)

# Show file sizes
sizes = []
export_size = os.path.getsize(export_file)
sizes.append(("jax.export", f"{export_size:,} bytes", f"{export_size / 1024:.0f} KB"))

sm_size = sum(
    os.path.getsize(os.path.join(dirpath, f))
    for dirpath, _, filenames in os.walk(savedmodel_base_dir)
    for f in filenames
)
sizes.append(("jax2tf SavedModel", f"{sm_size:,} bytes", f"{sm_size / 1024:.0f} KB"))

show_table(
    ["Export", "Size (bytes)", "Size (KB)"],
    sizes,
    title="Export file sizes",
    aligns=["left", "right", "right"],
)

আপনি প্রথমে সার্ভিং পাথ তুলনা সারণী এবং তারপরে এক্সপোর্ট ফাইলের আকারের সারণী দেখতে পাবেন, যেখানে jax.export ব্লব এবং SavedModel ডিরেক্টরির ডিস্ক-ভিত্তিক আকার তালিকাভুক্ত থাকবে।

মডেল আর্কিটেকচারের উপর ভিত্তি করে নয়, বরং আপনার ডেপ্লয়মেন্টের সীমাবদ্ধতার উপর ভিত্তি করে নির্বাচন করুন। যদি আপনার সার্ভিং প্রসেসে JAX আগে থেকেই চালু থাকে, তাহলে স্টার্টআপ ওয়ার্মআপ সহ jax.jit অথবা AOT কম্পাইলেশন হলো সবচেয়ে সংক্ষিপ্ত পথ। যদি আর্টিফ্যাক্টটিকে সেই প্রসেস থেকে বের হতে হয়, তাহলে jax.export আপনাকে JAX-এর ভেতরেই রাখে এবং jax2tf আপনাকে TensorFlow ইকোসিস্টেমের কাছে হস্তান্তর করে।

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

এই কোডল্যাবে যা কিছু লেখা হয়েছে, তা পডের /tmp ফোল্ডারে থাকে, যা পডের সাথে সাথে মুছে যায়। প্রথমে, ক্লাউড শেল থেকে এটি চালিয়ে আপনার প্রয়োজনীয় ফাইলগুলো কপি করে নিন:

kubectl cp jax-jupyter:/tmp/jax-course/exports ./jax-course-exports

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

kubectl delete -f deploy/jupyter.yaml

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

cd terraform
terraform destroy

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

gcloud container clusters list
gcloud compute instances list

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

১১. অভিনন্দন

আপনি নোটবুক থেকে একটি প্রশিক্ষিত ট্রান্সফরমার চেকপয়েন্ট নিয়ে সেটিকে একটি NVIDIA L4 GPU-তে চারটি পরিবেশন-উপযোগী আর্টিফ্যাক্টে পরিণত করেছেন।

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

• jax.jit হলো সবচেয়ে সহজ উপায় — র‍্যাপ করে কল করুন। প্রথম কলটি কম্পাইল হয়; পরবর্তী কলগুলো দ্রুত সম্পন্ন হয়। পাইথন-ভিত্তিক সার্ভিং (FastAPI, Flask)-এর জন্য এটি ভালো, যেখানে JAX আগে থেকেই ইনস্টল করা আছে।
• AOT কম্পাইলেশন ( lower() এবং তারপর compile() ) কম্পাইলেশনকে এক্সিকিউশন থেকে আলাদা করে। এটি স্টার্টআপের সময় একবার কম্পাইল করে, তারপর প্রথম-অনুরোধের কম্পাইলেশন বিলম্ব ছাড়াই পরিবেশন করে। lowered.as_text() ডিবাগিং এবং পারফরম্যান্স বিশ্লেষণের জন্য StableHLO IR দেখায়।
• jax.export একটি জিটেড ফাংশনকে একটি JAX-নেটিভ আর্টিফ্যাক্টে সিরিয়ালাইজ করে, যাতে StableHLO এবং কলিং মেটাডেটা থাকে। এর ফলে তৈরি হওয়া ফাইলটি মূল মডেল কোড ছাড়াই একটি সামঞ্জস্যপূর্ণ JAX রানটাইম দ্বারা লোড ও কল করা যায়। ডিফল্টরূপে এটি বর্তমান প্ল্যাটফর্মের জন্য এক্সপোর্ট করে; যখন আপনার একটি নির্দিষ্ট টার্গেটের প্রয়োজন হয়, তখন platforms=[...] ব্যবহার করুন।
• jax2tf ফাংশনটিকে একটি TensorFlow SavedModel-এ রূপান্তর করে। আপনার সার্ভিং পরিকাঠামো TensorFlow-ভিত্তিক হলে এটি ব্যবহার করুন। বর্তমান JAX রিলিজগুলিতে নেটিভ সিরিয়ালাইজেশনই ডিফল্ট; TensorFlow মডেলটি কোথায় এক্সিকিউট করবে, তার সাথে মিলিয়ে native_serialization_platforms সেট করুন।
• jax.ShapeDtypeStruct টার্গেট ব্যবহার করে কীভাবে একটি পুনর্নির্মিত আর্কিটেকচারে Orbax চেকপয়েন্ট পুনরুদ্ধার করা যায়, এবং jax.device_put ব্যবহার করে কীভাবে প্যারামিটারগুলিকে একটি একক ডিভাইসে স্থানান্তর করা যায়।
• কেন ফরওয়ার্ড-পাস টোকেন/সেকেন্ড অটোরেগ্রেসিভ জেনারেশন থ্রুপুট নয়, এবং কেন প্রতিটি নতুন ইনপুট শেপের জন্য নিজস্ব কম্পাইল করা এক্সিকিউটেবল প্রয়োজন

কোর্সের সারসংক্ষেপ

আটটি ল্যাবের মধ্যে আপনার জন্য রয়েছে:

1. L1–L3 : JAX সেট আপ করা, jit কম্পাইলেশন শেখা, এবং GPU এক্সিকিউশন প্রোফাইলিং করা।
2. L4 : একেবারে গোড়া থেকে একটি ট্রেনিং লুপ তৈরি করা — অপটিমাইজার, লস, গ্রেডিয়েন্ট আপডেট
3. L5 : মনোযোগের প্রক্রিয়াগুলো অন্বেষণ করা হয়েছে — সরল, SDPA, cuDNN সমন্বিত মনোযোগ
4. L6 : ডেটা প্যারালেলিজমের মাধ্যমে একাধিক GPU-তে স্কেলযোগ্য — মেশ, নেমডশার্ডিং, ডেটা-প্যারালাল শার্ডিং
5. L7 : সবকিছুকে একটি ট্রান্সফরমার ল্যাঙ্গুয়েজ মডেলে একত্রিত করেছে — Flax NNX, Orbax, generation
6. L8 : প্রোডাকশনের জন্য প্রশিক্ষিত মডেল প্রস্তুত করা হলো — JIT, AOT, jax.export , jax2tf

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

এখন আপনার কাছে জিপিইউ-তে JAX ব্যবহার করে বিল্ড করার মূল ওয়ার্কফ্লোটি রয়েছে। এর পরের সেরা পদক্ষেপ হলো এই অংশগুলোকে আরও বড় ও জটিল প্রজেক্টে পরিণত করা:

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

এখান থেকে মনোযোগ স্বতন্ত্র JAX বৈশিষ্ট্য শেখা থেকে সরে গিয়ে প্রকৌশলগত সিদ্ধান্ত নেওয়ার দিকে চলে যায়: কীভাবে মডেলগুলোকে দ্রুত, পুনরুৎপাদনযোগ্য, মেমরি-সাশ্রয়ী, ডিবাগযোগ্য এবং স্থাপনযোগ্য রাখা যায়।

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

• JAX AOT সংকলন
• jax.export গাইড
• jax2tf নেটিভ সিরিয়ালাইজেশনের জন্য JAX চেঞ্জলগ নোট
• ফ্ল্যাক্স এনএনএক্স এর প্রাথমিক বিষয়াবলী
• অরবাক্স চেকপয়েন্ট গাইড
• NGC-তে NVIDIA JAX কন্টেইনার