From 12eff26b6f4dafd2194893613dde1f928a076409 Mon Sep 17 00:00:00 2001 From: rajasekharporeddy Date: Tue, 28 Jul 2026 17:19:46 +0530 Subject: [PATCH] Fix the Attribute Error in gradient_accumulation_and_microbatching.ipynb example notebook --- ...dient_accumulation_and_microbatching.ipynb | 137 +++++++++--------- 1 file changed, 70 insertions(+), 67 deletions(-) diff --git a/examples/gradient_accumulation_and_microbatching.ipynb b/examples/gradient_accumulation_and_microbatching.ipynb index d1a33d19c..d8521cee5 100644 --- a/examples/gradient_accumulation_and_microbatching.ipynb +++ b/examples/gradient_accumulation_and_microbatching.ipynb @@ -1,10 +1,10 @@ { "cells": [ { + "cell_type": "markdown", "metadata": { "id": "-2Bmeh2sl_fp" }, - "cell_type": "markdown", "source": [ "# Summary\n", "\n", @@ -12,10 +12,12 @@ ] }, { + "cell_type": "code", + "execution_count": 1, "metadata": { "id": "5LUEGvWml-mH" }, - "cell_type": "code", + "outputs": [], "source": [ "# ! pip install -q \"optax @ git+https://github.com/google-deepmind/optax\"\n", "import jax\n", @@ -26,15 +28,13 @@ "import time\n", "from optax import microbatching\n", "import gc" - ], - "outputs": [], - "execution_count": 1 + ] }, { + "cell_type": "markdown", "metadata": { "id": "T5Ie-_tfqgAV" }, - "cell_type": "markdown", "source": [ "# Setup\n", "\n", @@ -43,10 +43,12 @@ ] }, { + "cell_type": "code", + "execution_count": 13, "metadata": { "id": "nInGNjjeqcvU" }, - "cell_type": "code", + "outputs": [], "source": [ "class TransformerBlock(nnx.Module):\n", " def __init__(self, hidden_size: int, num_heads: int, *, rngs: nnx.Rngs):\n", @@ -79,15 +81,15 @@ "\n", " logits = self.final_layer(x)\n", " return logits" - ], - "outputs": [], - "execution_count": 13 + ] }, { + "cell_type": "code", + "execution_count": 14, "metadata": { "id": "QEcf3p7fqoJ0" }, - "cell_type": "code", + "outputs": [], "source": [ "hidden_size = 512\n", "num_heads = 8\n", @@ -107,15 +109,15 @@ "\n", "adamw = optax.adamw(0.01)\n", "opt_state = adamw.init(params)" - ], - "outputs": [], - "execution_count": 14 + ] }, { + "cell_type": "code", + "execution_count": 15, "metadata": { "id": "iR9qzYp6qjDS" }, - "cell_type": "code", + "outputs": [], "source": [ "def loss_fn(params, batch):\n", " model = nnx.merge(graphdef, params)\n", @@ -133,15 +135,13 @@ " return params, opt_state\n", "\n", "update_fn.lower(params, opt_state, batch).compile()" - ], - "outputs": [], - "execution_count": 15 + ] }, { + "cell_type": "markdown", "metadata": { "id": "KlRFCRc5nHLP" }, - "cell_type": "markdown", "source": [ "# Part 1: Microbatching for Gradient Accumulation\n", "\n", @@ -149,10 +149,10 @@ ] }, { + "cell_type": "markdown", "metadata": { "id": "COxGVTPUq9bZ" }, - "cell_type": "markdown", "source": [ "### Option 1: Manual Gradient Accumulation\n", "\n", @@ -164,10 +164,12 @@ ] }, { + "cell_type": "code", + "execution_count": 8, "metadata": { "id": "SEf4ohj-q56F" }, - "cell_type": "code", + "outputs": [], "source": [ "@functools.partial(jax.jit, donate_argnums=(2,))\n", "def add_gradient(params, batch, accumulated_gradients):\n", @@ -184,15 +186,15 @@ "accumulated_gradients = optax.tree.zeros_like(params)\n", "add_gradient.lower(params, batch, accumulated_gradients).compile()\n", "update_params.lower(params, opt_state, accumulated_gradients).compile()" - ], - "outputs": [], - "execution_count": 8 + ] }, { + "cell_type": "code", + "execution_count": 9, "metadata": { "id": "FjIvz_kRrAXH" }, - "cell_type": "code", + "outputs": [], "source": [ "start_time = time.perf_counter()\n", "for i in range(accumulation_steps):\n", @@ -203,15 +205,13 @@ ")\n", "end_time = time.perf_counter()\n", "print('Total Time', end_time - start_time)" - ], - "outputs": [], - "execution_count": 9 + ] }, { + "cell_type": "markdown", "metadata": { "id": "gPfNKr5RrHmq" }, - "cell_type": "markdown", "source": [ "### Option 2: optax.MultiSteps\n", "\n", @@ -219,10 +219,12 @@ ] }, { + "cell_type": "code", + "execution_count": 11, "metadata": { "id": "vTMm84VyrDZV" }, - "cell_type": "code", + "outputs": [], "source": [ "multi_adam = optax.MultiSteps(adamw, accumulation_steps)\n", "\n", @@ -235,15 +237,15 @@ "\n", "multi_opt_state = multi_adam.init(params)\n", "update_fn_v2.lower(params, batch, multi_opt_state).compile()" - ], - "outputs": [], - "execution_count": 11 + ] }, { + "cell_type": "code", + "execution_count": 12, "metadata": { "id": "qi-zg9-5rGsg" }, - "cell_type": "code", + "outputs": [], "source": [ "start_time = time.perf_counter()\n", "for i in range(accumulation_steps):\n", @@ -252,15 +254,13 @@ "jax.block_until_ready((params, multi_opt_state))\n", "end_time = time.perf_counter()\n", "print('Total Time', end_time - start_time)" - ], - "outputs": [], - "execution_count": 12 + ] }, { + "cell_type": "markdown", "metadata": { "id": "SBGm3wY5rUIH" }, - "cell_type": "markdown", "source": [ "### Option 3: `microbatching.microbatch`\n", "\n", @@ -268,17 +268,19 @@ ] }, { + "cell_type": "markdown", "metadata": { "id": "I5HrKPuVrQiA" }, - "cell_type": "markdown", "source": [] }, { + "cell_type": "code", + "execution_count": 14, "metadata": { "id": "aBatfdC7rO--" }, - "cell_type": "code", + "outputs": [], "source": [ "@functools.partial(jax.jit, donate_argnums=(0, 2))\n", "def update_fn_v3(params, batch, opt_state):\n", @@ -293,29 +295,27 @@ "\n", "full_batch = jnp.vstack([batch]*accumulation_steps)\n", "update_fn_v3.lower(params, full_batch, opt_state).compile()" - ], - "outputs": [], - "execution_count": 14 + ] }, { + "cell_type": "code", + "execution_count": 17, "metadata": { "id": "JKxDiYD9rXMF" }, - "cell_type": "code", + "outputs": [], "source": [ "start_time = time.perf_counter()\n", "params, opt_state = jax.block_until_ready(update_fn_v3(params, full_batch, opt_state))\n", "end_time = time.perf_counter()\n", "print('Total Time', end_time - start_time)" - ], - "outputs": [], - "execution_count": 17 + ] }, { + "cell_type": "markdown", "metadata": { "id": "MZ2Xwu9znYYl" }, - "cell_type": "markdown", "source": [ "# Part 2: `microbatching.micro_vmap`\n", "\n", @@ -323,10 +323,12 @@ ] }, { + "cell_type": "code", + "execution_count": 12, "metadata": { "id": "fd4XmvYjnoh3" }, - "cell_type": "code", + "outputs": [], "source": [ "def expensive_function(x):\n", " return jax.nn.softmax(jnp.sin(jnp.outer(x, x))).sum(axis=0)\n", @@ -341,30 +343,28 @@ "gc.collect()\n", "print('Processed small batch', result.shape)\n", "print(result)" - ], - "outputs": [], - "execution_count": 12 + ] }, { + "cell_type": "code", + "execution_count": 11, "metadata": { "id": "nFGg2x5erilE" }, - "cell_type": "code", + "outputs": [], "source": [ "result = microbatching.micro_vmap(expensive_function, microbatch_size=32)(X)\n", "result = jax.block_until_ready(result)\n", "gc.collect()\n", "print('Processed Full Batch', result.shape)\n", "print(result)" - ], - "outputs": [], - "execution_count": 11 + ] }, { + "cell_type": "markdown", "metadata": { "id": "PWuT3BU1xJ2S" }, - "cell_type": "markdown", "source": [ "# Part 3: `microbatching.micro_grad`\n", "\n", @@ -372,10 +372,12 @@ ] }, { + "cell_type": "code", + "execution_count": 25, "metadata": { "id": "3hPNgICJuskf" }, - "cell_type": "code", + "outputs": [], "source": [ "def metrics_fn(per_example_grad):\n", " leaf_norms = jax.tree.map(jnp.linalg.norm, per_example_grad)\n", @@ -384,34 +386,35 @@ "grad_fn = microbatching.micro_grad(loss_fn, metrics_fn=metrics_fn, microbatch_size=8)\n", "\n", "grad, aux = jax.jit(grad_fn)(params, batch)" - ], - "outputs": [], - "execution_count": 25 + ] }, { + "cell_type": "code", + "execution_count": null, "metadata": { "id": "ph9T_WAWykls" }, - "cell_type": "code", + "outputs": [], "source": [ "# This shows the norm for the gradient of the embedding layer per example.\n", "# High uniformity of the norm values is encouraging.\n", - "aux.metrics['embedding']['embedding'].get_value()" - ], - "outputs": [], - "execution_count": 29 + "aux.metrics['embedding']['embedding'].value" + ] }, { + "cell_type": "markdown", "metadata": { "id": "XaMThtUpnakg" }, - "cell_type": "markdown", "source": [] } ], "metadata": { "colab": { "private_outputs": true + }, + "language_info": { + "name": "python" } }, "nbformat": 4,