jax

Bilingual

Flax debugging: making a hash of things
2 months ago
- The author was debugging a JAX/Flax NNX training loop where loss remained constant at 10.82, indicating the model was essentially guessing randomly.
- To verify whether model parameters were being updated, the author used hashing of parameter arrays by converting them to numpy bytes and computing a hash to detect any changes.
- The root cause was using @jax.jit instead of @nnx.jit, which prevented the in-place parameter updates required by Flax NNX. Switching to @nnx.jit resolved the issue, allowing the loss to decrease.
Writing an LLM from scratch, part 34a -- building a JAX training loop for an LLM training run
2 months ago
- The author built an LLM training loop from scratch using JAX, NNX, and Optax, based on notes from Sebastian Raschka's book.
- An 'A-to-A' model (identity mapping) was used to test the training loop before building a full LLM.
- Key challenges included OOM errors, JIT decorator issues, and learning how to use JAX for gradient accumulation.
- Gradient accumulation was implemented using Optax's MultiSteps wrapper, and gradient clipping used clip_by_global_norm.
- Checkpointing and restarting from checkpoints were enabled using Safetensors and Orbax, with a dispatcher in the learning rate extraction code.
- The training loop achieved 110k tokens/second throughput and a final loss of 0.001 on the A-to-A task.
- Learning rate scheduling used Optax's warmup_cosine_decay_schedule, with automatic charting of the learning rate over time.
- The next post will cover building and training the actual small LLM using this training harness.
JAX: commitment issues
2 months ago
- Creating an array with `jax.default_device` does not commit it to that device; JAX may move it elsewhere.
- Simple array lookups (e.g., `array[0]`) on an uncommitted CPU array can take over 5 seconds due to implicit GPU transfers.
- Explicitly committing an array with `jax.device_put(array, cpu0)` reduces lookup times dramatically (e.g., from 1.2s to <0.0002s).
- The issue stems from JAX moving the entire uncommitted array to the GPU for each lookup, causing severe performance degradation.
- For CPU-bound data like training datasets, always use `device_put` to nail arrays to the desired device and avoid hidden GPU overhead.
JAX backends and devices
2 months ago
- JAX defaults to loading data directly onto the GPU, leading to out-of-memory errors when the dataset exceeds VRAM capacity.
- The `jax.devices()` function only returns devices for the default backend (GPU if available, otherwise CPU) unless a specific backend is specified.
- To check available backends, JAX requires trying different backend options and catching errors, as there is no built-in listing method.
- JAX uses a default device concept, which can be temporarily changed using the `jax.default_device` context manager to load data onto the CPU instead of the GPU.
- The solution involves wrapping the data-loading code in `with jax.default_device(jax.devices("cpu")[0]):` to avoid GPU memory issues.
- Data loaded on the CPU can be moved to the GPU during training using `jax.device_put`.
Refs: Mutable Arrays in jax
2 months ago
- JAX Arrays are immutable, aiding optimized compilation but can limit expressiveness and performance for tasks like memory control.
- Refs are mutable arrays compatible with JAX transformations like jit and grad, allowing in-place reads and writes via indexing.
- Refs have constraints: they cannot be returned from jit functions, passed multiple times, or used in vmap with closed-over refs.
- Automatic differentiation works with pure functions using Refs internally, and ref arguments can be used for gradient accumulation via vjp.with_refs.
- Refs support in-place updates, enabling efficient gradient accumulation over microbatches without materializing separate arrays.
- foreach decorator uses Refs to write flexible loops with non-sequential access, avoiding scan's fixed pattern for better performance.
Defining new Jax types with hijax
2 months ago
- Hijax types (hi types) in JAX allow defining custom types beyond arrays for aggregate data with specific invariants and behaviors.
- To define a new hi type, subclass HiType, implement lo_ty, lower_val, and raise_val for lowering to arrays, and associate it with a value class via register_hitype.
- Primitives using hi types are defined as VJPHiPrimitive subclasses with type signatures mentioning the hi type, enforcing invariants through operations.
- Autodiff support requires implementing to_tangent_aval on the type and VJP/JVP rules on primitives, allowing custom tangent types (e.g., float arrays for quantized int8).
- vmap support involves implementing dec_rank and inc_rank on the type, defining a MappingSpec subclass, and batch rules on primitives, with explicit axis_size for hi type arguments.
- Sharding in explicit mode can be integrated by recording sharding data on the type (e.g., NamedSharding), consuming it in lo_ty, and propagating it in primitive typing rules.
- Example: quantized arrays (QArray) represent int8 values with per-row scales, defined as a hi type with primitives for quantization, dequantization, and operations like matmul.
- Hi types appear as single values in jaxprs (e.g., q8[...]), improving readability and ensuring invariants, unlike pytrees which expose individual leaves.
- Additional examples include rank-1 arrays (outer product representation) and generic tuples (TupTy), demonstrating flexibility in component structures and sharding delegation.
- Hi types support transformations like jit, grad, vmap, and scan, with custom handling for leading axes in scan via leading_axis_spec.

Flax debugging: making a hash of things

Writing an LLM from scratch, part 34a -- building a JAX training loop for an LLM training run

JAX: commitment issues

JAX backends and devices

Refs: Mutable Arrays in jax

Defining new Jax types with hijax