Skip to content

← Flax step 78 of 100

Medium Primitives

Parameter Counts and Ideal Shard Shares

Jump to the task

Before saving a model, it helps to know how many parameter values it contains. This exercise counts a Flax Dense layer's parameters and divides that count by a chosen number of shards. It does not create shards, write a checkpoint, or use distributed devices.

Count values, not bytes

For input shape (B, D), nn.Dense(F) creates:

  • A kernel of shape (D, F), containing D * F scalar values.
  • A bias of shape (F,), containing F scalar values.

The total is (D + 1) * F. The batch size B does not change it: every example uses the same kernel and bias. The seed changes the initialized kernel values, but not their count.

With D = 4 and F = 4, there are 16 kernel values and 4 bias values, for a total of 20. Dividing by two gives an ideal share of 10 values.

The return values are counts of scalar parameters, not storage bytes. For float32 values, the raw uncompressed parameter payload would use four bytes per value. A real checkpoint may also contain metadata, optimizer state and other arrays, and may use compression or a different storage dtype. Its file size need not equal this raw payload estimate.

An average is not a shard layout

Compute total / num_shards using ordinary division. The result may be fractional: six values divided among four hypothetical shards gives an average of 1.5 values. No shard literally contains half a scalar.

Actual array sharding must respect array dimensions and the selected partitioning scheme. Some values may be replicated rather than split. A quotient alone cannot establish that an even partition is possible, how many values each device stores, or how many files will be written.

Where real checkpointing differs

Distributed checkpoint systems coordinate stored array data and metadata. Device shards, replicas and storage chunks are different concepts. Orbax can use TensorStore-backed formats such as OCDBT, whose files do not follow a universal one-file-per-device-shard layout.

Saving across more hosts does not guarantee proportional speedup, and partitioning data does not itself provide redundancy if a shard is lost. Restore also needs compatible structure, shapes, dtypes and placement. Those concerns require real save/restore tests; none is simulated by the two numbers returned here. This is also separate from gradient checkpointing, which trades derivative storage for recomputation.

Problem

Implement distributed_checkpoint_simulate(seed, x, num_shards, features):

  1. Convert seed, width and shard count to Python integers.
  2. Initialize default nn.Dense(features) with jax.random.PRNGKey(int(seed)) and x.
  3. Read the kernel and bias from the returned params collection.
  4. Add their .size values to obtain total_params.
  5. Compute ideal_share = total_params / num_shards, without rounding.
  6. Return a JAX array of shape (2,): [total_params, ideal_share].

x is a finite float32 JAX array of shape (B, D), with B, D >= 1. Width and shard count are positive integer-valued scalars, and seed is a nonnegative integer-valued scalar. They may arrive as floats. Keep the default bias enabled and leave x unchanged.

The result is floating point so it can represent a fractional share. Do not multiply by bytes per value, use integer division, pad the count, or write files. Output tests verify the counts; they cannot prove that a submission initialized a Dense layer instead of using the formula.