[JAX] Add tutorial for integrating TE/JAX quantization into an existing framework (#2423)

* Tutorial for integration te/jax quantization into an existing framework

Signed-off-by: Jeremy Berchtold <[email protected]>

* add todos

Signed-off-by: Jeremy Berchtold <[email protected]>

* support nvfp4 sr rng key, move wrapper module into TE itself, fix bfloat16 cast

Signed-off-by: Jeremy Berchtold <[email protected]>

* update docstrings

Signed-off-by: Jeremy Berchtold <[email protected]>

* Fix QKV proj and out proj in Flax example transformer layer

Signed-off-by: Jeremy Berchtold <[email protected]>

* Use fused attention in quickstart_jax example

Signed-off-by: Jeremy Berchtold <[email protected]>

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

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

* remat policy

Signed-off-by: Jeremy Berchtold <[email protected]>

* add tutorial to docs

Signed-off-by: Jeremy Berchtold <[email protected]>

* update title

Signed-off-by: Jeremy Berchtold <[email protected]>

* remove unused dtype from TE DPA module

Signed-off-by: Jeremy Berchtold <[email protected]>

* Fix notebook title

Signed-off-by: Jeremy Berchtold <[email protected]>

* Fix lint

Signed-off-by: Jeremy Berchtold <[email protected]>

* Add explanation of flax module wrapper

Signed-off-by: Jeremy Berchtold <[email protected]>

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

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

---------

Signed-off-by: Jeremy Berchtold <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: jberchtold-nvidia

docs/examples/quickstart_jax.ipynb

@@ -19,12 +19,12 @@ def speedometer(
     variables: Any,
     input: jnp.ndarray,
     output_grad: jnp.ndarray,
-    dropout_key: jax.random.PRNGKey,
     model_init_fn: Callable = None,
     forward_kwargs: dict = {},
     autocast_kwargs: Optional[dict] = None,
     timing_iters: int = 50,
     warmup_iters: int = 50,
+    rngs: Dict[str, jax.random.PRNGKey] = None,
 ) -> None:
     """Measure average runtime for a JAX module
     Perform forward and backward passes .
@@ -33,19 +33,21 @@ def speedometer(
         autocast_kwargs = {"enabled": False}
         model_init_fn = None
 
+    if rngs is None:
+        rngs = {}
+
     train_step_fn = create_train_step_fn(model_apply_fn, autocast_kwargs, forward_kwargs)
 
     # Warm up runs
-    key = dropout_key
     for _ in range(warmup_iters):
-        key, step_key = jax.random.split(key)
-        loss, (param_grads, other_grads) = train_step_fn(variables, input, output_grad, step_key)
+        rngs, step_rngs = _split_step_rngs(rngs)
+        loss, (param_grads, other_grads) = train_step_fn(variables, input, output_grad, step_rngs)
 
     # Timing runs
     start = time.time()
     for _ in range(timing_iters):
-        key, step_key = jax.random.split(key)
-        loss, (param_grads, other_grads) = train_step_fn(variables, input, output_grad, step_key)
+        rngs, step_rngs = _split_step_rngs(rngs)
+        loss, (param_grads, other_grads) = train_step_fn(variables, input, output_grad, step_rngs)
     end = time.time()
 
     print(f"Mean time: {(end - start) * 1000 / timing_iters} ms")
@@ -63,8 +65,12 @@ def create_train_step_fn(
     if forward_kwargs is None:
         forward_kwargs = {}
 
-    def loss_fn(variables: Any, inp: jnp.ndarray, grad_target: jnp.ndarray, dropout_key):
-        rngs = {"dropout": dropout_key}
+    def loss_fn(
+        variables: Any,
+        inp: jnp.ndarray,
+        grad_target: jnp.ndarray,
+        rngs: Dict[str, jax.random.PRNGKey],
+    ):
         with te.autocast(**autocast_kwargs):
             # Forward Pass: Apply the model using current parameters and variables
             call_kwargs = {**forward_kwargs, "rngs": rngs}
@@ -84,3 +90,16 @@ def fwd_bwd_fn(*args, **kwargs):
 
     # JIT-compile the fwd_bwd_fn
     return jax.jit(fwd_bwd_fn)
+
+
+def _split_step_rngs(
+    rngs: Dict[str, jax.random.PRNGKey],
+) -> Tuple[Dict[str, jax.random.PRNGKey], Dict[str, jax.random.PRNGKey]]:
+    """Splits each RNG in the rngs dictionary for a new step."""
+    step_rngs = {}
+    new_rngs = {}
+    for name, key in rngs.items():
+        new_key, step_key = jax.random.split(key)
+        new_rngs[name] = new_key
+        step_rngs[name] = step_key
+    return new_rngs, step_rngs