Remove weight sharing between different iterations of the transformerLayer

Signed-off-by: tdophung <tdophung@nvidia.com>

[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Signed-off-by: tdophung <tdophung@nvidia.com>

Author: tdophung

docs/examples/quickstart_jax.ipynb

@@ -51,7 +51,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 25,
    "id": "881fd001",
    "metadata": {},
    "outputs": [],
@@ -65,7 +65,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 8,
+   "execution_count": 26,
    "id": "d5284a38",
    "metadata": {},
    "outputs": [],
@@ -79,7 +79,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": 27,
    "id": "a4d1cfdc",
    "metadata": {},
    "outputs": [],
@@ -173,7 +173,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 28,
    "id": "8b44649d",
    "metadata": {},
    "outputs": [],
@@ -194,7 +194,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 29,
    "id": "e44ed26d",
    "metadata": {},
    "outputs": [
@@ -224,7 +224,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 30,
    "id": "de91af7a",
    "metadata": {},
    "outputs": [
@@ -250,15 +250,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 31,
    "id": "037bc8d9",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Mean time: 27.269372940063477 ms\n"
+      "Mean time: 28.229827880859375 ms\n"
      ]
     }
    ],
@@ -308,7 +308,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 32,
    "id": "bed20d6b",
    "metadata": {},
    "outputs": [],
@@ -328,7 +328,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 33,
    "id": "56105579",
    "metadata": {},
    "outputs": [],
@@ -424,16 +424,15 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 34,
    "id": "5146cd99",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Basic TE parameter shapes: {'BasicTEMLP_0': {'DenseGeneral_0': {'bias': LogicallyPartitioned(value=(16384,), names=(), mesh=None, rules=None), 'kernel': LogicallyPartitioned(value=(4096, 16384), names=(), mesh=None, rules=None)}, 'DenseGeneral_1': {'bias': LogicallyPartitioned(value=(4096,), names=(), mesh=None, rules=None), 'kernel': LogicallyPartitioned(value=(16384, 4096), names=(), mesh=None, rules=None)}}, 'DenseGeneral_0': {'bias': LogicallyPartitioned(value=(12288,), names=(), mesh=None, rules=None), 'kernel': LogicallyPartitioned(value=(4096, 12288), names=(), mesh=None, rules=None)}, 'DenseGeneral_1': {'bias': LogicallyPartitioned(value=(4096,), names=(), mesh=None, rules=None), 'kernel': LogicallyPartitioned(value=(4096, 4096), names=(), mesh=None, rules=None)}, 'LayerNorm_0': {'ln_bias': LogicallyPartitioned(value=(4096,), names=('embed',), mesh=None, rules=None), 'scale': LogicallyPartitioned(value=(4096,), names=('embed',), mesh=None, rules=None)}, 'LayerNorm_1': {'ln_bias': LogicallyPartitioned(value=(4096,), names=('embed',), mesh=None, rules=None), 'scale': LogicallyPartitioned(value=(4096,), names=('embed',), mesh=None, rules=None)}}\n",
-      "Mean time: 17.397570610046387 ms\n"
+      "Mean time: 17.390952110290527 ms\n"
      ]
     }
    ],
@@ -449,21 +448,12 @@
     "\n",
     "te_params = basic_te_transformer.init(key, x, attention_mask=None, deterministic=False)\n",
     "\n",
-    "# Extract the 'params' pytrees\n",
-    "basic_params = params['params']\n",
-    "te_params_template = te_params['params']\n",
-    "\n",
-    "print(f\"Basic TE parameter shapes: {jax.tree_util.tree_map(lambda x: x.shape, te_params_template)}\")\n",
-    "\n",
-    "shared_te_params = {}\n",
-    "shared_te_params['params'] = utils.share_parameters_with_basic_te_model(basic_params, te_params_template)\n",
-    "\n",
     "# Test forward pass\n",
     "y = basic_te_transformer.apply(te_params, x, attention_mask=None, deterministic=True)\n",
     "\n",
     "utils.speedometer(\n",
     "    model_apply_fn=basic_te_transformer.apply,\n",
-    "    variables=shared_te_params,  # Ensure the correct `params` is passed\n",
+    "    variables=te_params,  # Ensure the correct `params` is passed\n",
     "    input=x,\n",
     "    output_grad=dy,\n",
     "    dropout_key=dropout_key,\n",
@@ -502,7 +492,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 35,
    "id": "11203785",
    "metadata": {},
    "outputs": [],
@@ -570,58 +560,41 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 18,
+   "execution_count": 36,
    "id": "114de14f",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Fused TE parameter shapes: {'DenseGeneral_0': {'bias': LogicallyPartitioned(value=(4096,), names=(), mesh=None, rules=None), 'kernel': LogicallyPartitioned(value=(4096, 4096), names=(), mesh=None, rules=None)}, 'LayerNormDenseGeneral_0': {'bias': LogicallyPartitioned(value=(12288,), names=(), mesh=None, rules=None), 'kernel': LogicallyPartitioned(value=(4096, 12288), names=(), mesh=None, rules=None), 'ln_bias': LogicallyPartitioned(value=(4096,), names=('embed',), mesh=None, rules=None), 'scale': LogicallyPartitioned(value=(4096,), names=('embed',), mesh=None, rules=None)}, 'LayerNormMLP_0': {'ln_bias': LogicallyPartitioned(value=(4096,), names=('embed',), mesh=None, rules=None), 'scale': LogicallyPartitioned(value=(4096,), names=('embed',), mesh=None, rules=None), 'wi_bias': LogicallyPartitioned(value=(1, 16384), names=('act', 'mlp'), mesh=None, rules=None), 'wi_kernel': LogicallyPartitioned(value=(4096, 1, 16384), names=('embed', 'act', 'mlp'), mesh=None, rules=None), 'wo_bias': LogicallyPartitioned(value=(4096,), names=('embed',), mesh=None, rules=None), 'wo_kernel': LogicallyPartitioned(value=(16384, 4096), names=('mlp', 'embed'), mesh=None, rules=None)}}\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
-    "import quickstart_jax_utils\n",
-    "importlib.reload(quickstart_jax_utils)\n",
-    "\n",
     "fused_te_transformer = FusedTETransformerLayer(\n",
     "    hidden_size, \n",
     "    ffn_hidden_size, \n",
     "    num_attention_heads\n",
     ")\n",
     "\n",
-    "fused_te_params = fused_te_transformer.init(key, x, attention_mask=None, deterministic=False)\n",
-    "\n",
-    "fused_te_params_template = fused_te_params['params']\n",
-    "print(f\"Fused TE parameter shapes: {jax.tree_util.tree_map(lambda x: x.shape, fused_te_params_template)}\")"
+    "fused_te_params = fused_te_transformer.init(key, x, attention_mask=None, deterministic=False)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 19,
+   "execution_count": 37,
    "id": "6b0c705e",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Mean time: 18.0991792678833 ms\n"
+      "Mean time: 18.087706565856934 ms\n"
      ]
     }
    ],
    "source": [
-    "shared_fused_te_params = {}\n",
-    "shared_fused_te_params['params'] = utils.share_fused_parameters_with_basic_te_model(basic_params, fused_te_params_template)\n",
-    "\n",
     "# Test forward pass\n",
     "y = fused_te_transformer.apply(fused_te_params, x, attention_mask=None, deterministic=True)\n",
     "\n",
     "utils.speedometer(\n",
     "    model_apply_fn=fused_te_transformer.apply,\n",
-    "    variables=shared_fused_te_params,  # Ensure the correct `params` is passed\n",
+    "    variables=fused_te_params,  # Ensure the correct `params` is passed\n",
     "    input=x,\n",
     "    output_grad=dy,\n",
     "    dropout_key=dropout_key,\n",
@@ -639,18 +612,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 20,
+   "execution_count": 38,
    "id": "7496b159",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "TE TransformerLayer parameter shapes: {'attention': {'out': {'bias': LogicallyPartitioned(value=(4096,), names=('nvte_w_no_shard',), mesh=None, rules=None), 'kernel': LogicallyPartitioned(value=(4096, 4096), names=('nvte_w_tp', 'nvte_w_fsdp'), mesh=None, rules=None)}, 'qkv': {'bias': LogicallyPartitioned(value=(3, 4096), names=('nvte_w_joined', 'nvte_w_tp'), mesh=None, rules=None), 'kernel': LogicallyPartitioned(value=(4096, 3, 4096), names=('nvte_w_fsdp', 'nvte_w_joined', 'nvte_w_tp'), mesh=None, rules=None), 'ln_bias': LogicallyPartitioned(value=(4096,), names=('nvte_w_no_shard',), mesh=None, rules=None), 'scale': LogicallyPartitioned(value=(4096,), names=('nvte_w_no_shard',), mesh=None, rules=None)}}, 'mlp': {'ln_bias': LogicallyPartitioned(value=(4096,), names=('nvte_w_no_shard',), mesh=None, rules=None), 'scale': LogicallyPartitioned(value=(4096,), names=('nvte_w_no_shard',), mesh=None, rules=None), 'wi_bias': LogicallyPartitioned(value=(1, 16384), names=('nvte_w_joined', 'nvte_w_tp'), mesh=None, rules=None), 'wi_kernel': LogicallyPartitioned(value=(4096, 1, 16384), names=('nvte_w_fsdp', 'nvte_w_joined', 'nvte_w_tp'), mesh=None, rules=None), 'wo_bias': LogicallyPartitioned(value=(4096,), names=('nvte_w_no_shard',), mesh=None, rules=None), 'wo_kernel': LogicallyPartitioned(value=(16384, 4096), names=('nvte_w_tp', 'nvte_w_fsdp'), mesh=None, rules=None)}, 'relpos_bias': {'rel_embedding': LogicallyPartitioned(value=(32, 32), names=('heads', 'relpos_buckets'), mesh=None, rules=None)}}\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "te_transformer = te_flax.TransformerLayer(\n",
     "    hidden_size=hidden_size,\n",
@@ -662,39 +627,30 @@
     "    use_bias=True\n",
     "    )\n",
     "\n",
-    "te_transformer_params = te_transformer.init(key, x, deterministic=False)\n",
-    "\n",
-    "te_transformer_params_template = te_transformer_params['params']\n",
-    "print(f\"TE TransformerLayer parameter shapes: {jax.tree_util.tree_map(lambda x: x.shape, te_transformer_params_template)}\")"
+    "te_transformer_params = te_transformer.init(key, x, deterministic=False)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": 39,
    "id": "6ec0f60e",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Mean time: 11.84274673461914 ms\n"
+      "Mean time: 12.37576961517334 ms\n"
      ]
     }
    ],
    "source": [
-    "import quickstart_jax_utils\n",
-    "importlib.reload(quickstart_jax_utils)\n",
-    "\n",
-    "shared_te_transformer_params = {}\n",
-    "shared_te_transformer_params['params'] = utils.share_parameters_with_transformerlayer_te_model(basic_params, te_transformer_params_template)\n",
-    "\n",
     "# Test forward pass\n",
     "y = te_transformer.apply(te_transformer_params, x, attention_mask=None, deterministic=True)\n",
     "\n",
     "utils.speedometer(\n",
     "    model_apply_fn=te_transformer.apply,\n",
-    "    variables=shared_te_transformer_params,  # Ensure the correct `params` is passed\n",
+    "    variables=te_transformer_params,  # Ensure the correct `params` is passed\n",
     "    input=x,\n",
     "    output_grad=dy,\n",
     "    dropout_key=dropout_key,\n",
@@ -730,18 +686,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 22,
+   "execution_count": 40,
    "id": "b2aaa8ef",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "TE TransformerLayer vars: {'fp8_metas': {'attention': {'out': {'grad_amax_history': (16,), 'grad_scale': (1,), 'kernel_amax_history': (16,), 'kernel_scale': (1,), 'x_amax_history': (16,), 'x_scale': (1,)}, 'qkv': {'grad_amax_history': (16,), 'grad_scale': (1,), 'kernel_amax_history': (16,), 'kernel_scale': (1,), 'x_amax_history': (16,), 'x_scale': (1,)}}, 'mlp': {'grad_0_amax_history': (16,), 'grad_0_scale': (1,), 'grad_1_amax_history': (16,), 'grad_1_scale': (1,), 'kernel_0_amax_history': (16,), 'kernel_0_scale': (1,), 'kernel_1_amax_history': (16,), 'kernel_1_scale': (1,), 'x_0_amax_history': (16,), 'x_0_scale': (1,), 'x_1_amax_history': (16,), 'x_1_scale': (1,)}}, 'params': {'attention': {'out': {'bias': LogicallyPartitioned(value=(4096,), names=('nvte_w_no_shard',), mesh=None, rules=None), 'kernel': LogicallyPartitioned(value=(4096, 4096), names=('nvte_w_tp', 'nvte_w_fsdp'), mesh=None, rules=None)}, 'qkv': {'bias': LogicallyPartitioned(value=(3, 4096), names=('nvte_w_joined', 'nvte_w_tp'), mesh=None, rules=None), 'kernel': LogicallyPartitioned(value=(4096, 3, 4096), names=('nvte_w_fsdp', 'nvte_w_joined', 'nvte_w_tp'), mesh=None, rules=None), 'ln_bias': LogicallyPartitioned(value=(4096,), names=('nvte_w_no_shard',), mesh=None, rules=None), 'scale': LogicallyPartitioned(value=(4096,), names=('nvte_w_no_shard',), mesh=None, rules=None)}}, 'mlp': {'ln_bias': LogicallyPartitioned(value=(4096,), names=('nvte_w_no_shard',), mesh=None, rules=None), 'scale': LogicallyPartitioned(value=(4096,), names=('nvte_w_no_shard',), mesh=None, rules=None), 'wi_bias': LogicallyPartitioned(value=(1, 16384), names=('nvte_w_joined', 'nvte_w_tp'), mesh=None, rules=None), 'wi_kernel': LogicallyPartitioned(value=(4096, 1, 16384), names=('nvte_w_fsdp', 'nvte_w_joined', 'nvte_w_tp'), mesh=None, rules=None), 'wo_bias': LogicallyPartitioned(value=(4096,), names=('nvte_w_no_shard',), mesh=None, rules=None), 'wo_kernel': LogicallyPartitioned(value=(16384, 4096), names=('nvte_w_tp', 'nvte_w_fsdp'), mesh=None, rules=None)}, 'relpos_bias': {'rel_embedding': LogicallyPartitioned(value=(32, 32), names=('heads', 'relpos_buckets'), mesh=None, rules=None)}}}\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "from transformer_engine.common.recipe import Format, DelayedScaling\n",
     "\n",
@@ -760,40 +708,28 @@
     "\n",
     "with te.fp8_autocast(enabled=True, fp8_recipe=fp8_recipe):\n",
     "    te_transformer_params = te_transformer.init(key, x, deterministic=False)\n",
-    "    \n",
-    "    # When using FP8, we need to preserve the fp8_metas collection\n",
-    "    # that was created during initialization within the fp8_autocast context.\n",
-    "    # Only the 'params' are shared from basic_params, but fp8_metas must come from\n",
-    "    # the FP8-initialized model.\n",
-    "    shared_te_transformer_params = {}\n",
-    "    shared_te_transformer_params['params'] = utils.share_parameters_with_transformerlayer_te_model(basic_params, te_transformer_params_template)\n",
-    "    print(f\"TE TransformerLayer vars: {jax.tree_util.tree_map(lambda x: x.shape, te_transformer_params)}\")\n",
-    "\n",
-    "    if 'fp8_metas' in te_transformer_params:\n",
-    "        shared_te_transformer_params['fp8_metas'] = te_transformer_params['fp8_metas']\n",
-    "\n",
     "    y = te_transformer.apply(te_transformer_params, x, attention_mask=None, deterministic=True)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 23,
+   "execution_count": 41,
    "id": "b9cdbf22",
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Mean time: 7.96757698059082 ms\n"
+      "Mean time: 7.956786155700684 ms\n"
      ]
     }
    ],
    "source": [
     "utils.speedometer(\n",
     "    model_apply_fn=te_transformer.apply,\n",
     "    model_init_fn=te_transformer.init,\n",
-    "    variables=shared_te_transformer_params,  # Includes both params and fp8_metas\n",
+    "    variables=te_transformer_params,  # Includes both params and fp8_metas\n",
     "    input=x,\n",
     "    output_grad=dy,\n",
     "    dropout_key=dropout_key,\n",
@@ -808,6 +744,18 @@
    "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.3"
   }
  },
  "nbformat": 4,