Module 9. Pretraining: Data, Scaling Laws, and Optimizers

import hashlib
import importlib.util
import json
import math
import subprocess
import sys
from dataclasses import dataclass
from typing import Sequence

if importlib.util.find_spec("numpy") is None:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "numpy>=1.24,<3"])

import numpy as np

np.set_printoptions(precision=6, suppress=True)
REFERENCE_JSON = '{"break_even":{"break_even_generated_tokens":2123682119443.3926,"extra_inference_flops_per_token_for_large_model":123940000000.0,"large_model_parameters":70000000000.0,"large_training_flops":4.59490838116186e+23,"large_training_tokens":1094025805038.538,"small_training_flops":7.227e+23,"target_loss":1.9483096526156487,"training_flops_saved_by_large_model":2.6320916188381408e+23},"data_mixture":{"allocations":[{"allocated_tokens":8250000000000.001,"effective_epochs":1.3750000000000002,"name":"web","normalized_weight":0.55,"unique_tokens":6000000000000.0},{"allocated_tokens":3000000000000.0,"effective_epochs":3.0,"name":"code","normalized_weight":0.2,"unique_tokens":1000000000000.0},{"allocated_tokens":2250000000000.0,"effective_epochs":5.625,"name":"math","normalized_weight":0.15,"unique_tokens":400000000000.0},{"allocated_tokens":1500000000000.0,"effective_epochs":2.5,"name":"books","normalized_weight":0.1,"unique_tokens":600000000000.0}],"slices":[{"name":"web","sampling_weight":0.55,"unique_tokens":6000000000000.0},{"name":"code","sampling_weight":0.2,"unique_tokens":1000000000000.0},{"name":"math","sampling_weight":0.15,"unique_tokens":400000000000.0},{"name":"books","sampling_weight":0.1,"unique_tokens":600000000000.0}],"total_training_tokens":15000000000000.0},"muon":{"after_cubic_tail":{"condition_number":1.0010334534088043,"maximum_singular_value":0.9999999999999993,"minimum_singular_value":0.9989676135145276,"orthogonality_error_frobenius":0.004944701975528774},"after_quintic":{"condition_number":1.6612800631385443,"maximum_singular_value":1.132919267754335,"minimum_singular_value":0.6819556153668558,"orthogonality_error_frobenius":2.164423918684206},"before":{"condition_number":5.12321315773202,"maximum_singular_value":13.807749133977339,"minimum_singular_value":2.695134617449696,"orthogonality_error_frobenius":438.60544001408556},"cubic_steps":3,"quintic_steps":5,"seed":7,"shape":[64,32],"wide_matches_transpose":true,"wide_shape":[32,64]},"optimal_allocations":[{"compute_flops":7.23e+23,"model_parameters":35669778435.426346,"predicted_loss":1.922490503187668,"tokens_per_parameter":94.70790361664052,"training_tokens":3378209938089.281},{"compute_flops":2.6e+24,"model_parameters":63580317550.129486,"predicted_loss":1.881010676151493,"tokens_per_parameter":107.19553898543344,"training_tokens":6815526408651.144},{"compute_flops":1e+25,"model_parameters":116822979225.12761,"predicted_loss":1.8453201291884438,"tokens_per_parameter":122.12151819239831,"training_tokens":14266599582731.592},{"compute_flops":1e+26,"model_parameters":330476955191.908,"predicted_loss":1.7990634812441677,"tokens_per_parameter":152.60416548167666,"training_tokens":50432159957986.57}],"pipeline":{"minimum_microbatches_for_5pct":134,"rows":[{"bubble_fraction":0.1794871794871795,"microbatches":32,"stages":8},{"bubble_fraction":0.05185185185185185,"microbatches":128,"stages":8}]},"qk_clip":{"clipped_view_after_regrowth":[59.89500000000002,100.0,100.0,100.0],"corrected_max_logits":[45.0,100.00000000000001,100.00000000000001,80.0],"gamma":[1.0,0.8333333333333334,0.38461538461538464,1.0],"key_scale":[1.0,0.9128709291752769,0.6201736729460423,1.0],"max_logits":[45.0,120.0,260.0,80.0],"query_scale":[1.0,0.9128709291752769,0.6201736729460423,1.0],"reapplied_corrected_max_logits":[59.89500000000002,100.0,100.0,99.99999999999999],"regrown_after_three_10pct_steps":[59.89500000000002,133.10000000000005,133.10000000000005,106.48000000000003],"split":0.5,"threshold":100.0},"scaling_law":{"compute_multiplier":6.0,"data_coefficient":410.7,"data_exponent":0.28,"irreducible_loss":1.69,"model_coefficient":406.4,"model_exponent":0.34},"schema_version":1,"small_model_scenario":{"loss_gap_to_same_budget_optimum":0.025804333208937935,"model_parameters":8030000000.0,"predicted_loss":1.9483096526156487,"same_budget_optimum":{"compute_flops":7.227e+23,"model_parameters":35663093470.54878,"predicted_loss":1.9225053194067108,"tokens_per_parameter":94.70409988264016,"training_tokens":3377441166158.784},"tokens_per_parameter":1867.9950186799501,"training_flops":7.227e+23,"training_tokens":15000000000000.0},"training_run":{"accelerator_count":2048,"effective_cluster_flops_per_second":8.6016e+17,"estimated_days":9.72445048983135,"estimated_seconds":840192.5223214286,"model_flop_utilization":0.42,"model_parameters":8030000000.0,"peak_flops_per_accelerator":1000000000000000.0,"total_training_flops":7.227e+23,"training_tokens":15000000000000.0,"training_tokens_per_second":17853051.05853051},"training_state_memory":{"parameter_count":7000000000,"policy":{"first_moment_bytes":4,"gradient_bytes":2,"master_weight_bytes":4,"model_weight_bytes":2,"second_moment_bytes":4,"total_bytes_per_parameter":16},"rows":[{"bytes_per_rank":112000000000.0,"decimal_gb_per_rank":112.0,"gib_per_rank":104.3081283569336,"zero_stage":0},{"bytes_per_rank":29312500000.0,"decimal_gb_per_rank":29.3125,"gib_per_rank":27.299392968416214,"zero_stage":1},{"bytes_per_rank":15531250000.0,"decimal_gb_per_rank":15.53125,"gib_per_rank":14.46460373699665,"zero_stage":2},{"bytes_per_rank":1750000000.0,"decimal_gb_per_rank":1.75,"gib_per_rank":1.6298145055770874,"zero_stage":3}],"world_size":64}}'
REFERENCE_SHA256 = "224ceb35b077fa5ad4081ef7c4c77607599c53694b52bc3f13d40fd739a86331"
REFERENCE = json.loads(REFERENCE_JSON)
canonical = json.dumps(REFERENCE, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
assert hashlib.sha256(canonical.encode("utf-8")).hexdigest() == REFERENCE_SHA256
assert REFERENCE["schema_version"] == 1
print('Embedded M9 contract verified' + ":", REFERENCE_SHA256[:12])

def _real(name: str, value: float) -> float:
    """Return a finite, non-boolean real number."""

    if isinstance(value, bool) or not isinstance(value, (int, float, np.integer, np.floating)):
        raise ValueError(f"{name} must be a real number")
    number = float(value)
    if not math.isfinite(number):
        raise ValueError(f"{name} must be finite")
    return number


def _positive_finite(name: str, value: float) -> float:
    """Return a strictly positive finite real number."""

    number = _real(name, value)
    if number <= 0.0:
        raise ValueError(f"{name} must be positive")
    return number


def _non_negative_finite(name: str, value: float) -> float:
    """Return a non-negative finite real number."""

    number = _real(name, value)
    if number < 0.0:
        raise ValueError(f"{name} must be non-negative")
    return number


def _positive_int(name: str, value: int) -> int:
    """Return a strictly positive, non-boolean integer."""

    if isinstance(value, bool) or not isinstance(value, (int, np.integer)) or int(value) <= 0:
        raise ValueError(f"{name} must be a positive integer")
    return int(value)


def _non_negative_int(name: str, value: int) -> int:
    """Return a non-negative, non-boolean integer."""

    if isinstance(value, bool) or not isinstance(value, (int, np.integer)) or int(value) < 0:
        raise ValueError(f"{name} must be a non-negative integer")
    return int(value)


def _fraction(name: str, value: float, *, allow_zero: bool = True) -> float:
    """Return a finite fraction in ``[0, 1]`` or ``(0, 1]``."""

    number = _real(name, value)
    lower_ok = number >= 0.0 if allow_zero else number > 0.0
    if not lower_ok or number > 1.0:
        interval = "[0, 1]" if allow_zero else "(0, 1]"
        raise ValueError(f"{name} must lie in {interval}")
    return number


def bytes_to_gib(byte_count: float) -> float:
    """Convert a finite non-negative byte count to binary GiB."""

    return _non_negative_finite("byte_count", byte_count) / 2**30


def bytes_to_decimal_gb(byte_count: float) -> float:
    """Convert a finite non-negative byte count to decimal GB."""

    return _non_negative_finite("byte_count", byte_count) / 10**9

@dataclass(frozen=True)
class ComputeAllocation:
    """A model/data allocation under a fixed training-compute budget."""

    compute_flops: float
    model_parameters: float
    training_tokens: float
    predicted_loss: float

    @property
    def tokens_per_parameter(self) -> float:
        """Return the data-to-model ratio for this allocation."""

        return self.training_tokens / self.model_parameters


@dataclass(frozen=True)
class ScalingLaw:
    """Two-term Chinchilla-style loss fit with an explicit compute convention.

    ``compute_multiplier`` is the declared approximation in
    ``training_flops = compute_multiplier * N * D``.  It is a modelling convention,
    not a measured wall-clock quantity.
    """

    irreducible_loss: float
    model_coefficient: float
    data_coefficient: float
    model_exponent: float
    data_exponent: float
    compute_multiplier: float = 6.0

    def __post_init__(self) -> None:
        _non_negative_finite("irreducible_loss", self.irreducible_loss)
        _positive_finite("model_coefficient", self.model_coefficient)
        _positive_finite("data_coefficient", self.data_coefficient)
        _positive_finite("model_exponent", self.model_exponent)
        _positive_finite("data_exponent", self.data_exponent)
        _positive_finite("compute_multiplier", self.compute_multiplier)

    def loss(self, model_parameters: float, training_tokens: float) -> float:
        """Return the fitted loss for positive ``N`` and ``D``."""

        n = _positive_finite("model_parameters", model_parameters)
        d = _positive_finite("training_tokens", training_tokens)
        return (
            self.irreducible_loss
            + self.model_coefficient * n ** (-self.model_exponent)
            + self.data_coefficient * d ** (-self.data_exponent)
        )

    def training_flops(self, model_parameters: float, training_tokens: float) -> float:
        """Return the declared arithmetic training-FLOP estimate ``kappa*N*D``."""

        n = _positive_finite("model_parameters", model_parameters)
        d = _positive_finite("training_tokens", training_tokens)
        return self.compute_multiplier * n * d

    def optimal_split(self, compute_flops: float) -> ComputeAllocation:
        """Return the analytical compute-optimal ``N`` and ``D`` for this fit."""

        c = _positive_finite("compute_flops", compute_flops)
        alpha = self.model_exponent
        beta = self.data_exponent
        n = (
            (alpha * self.model_coefficient)
            / (beta * self.data_coefficient)
            * (c / self.compute_multiplier) ** beta
        ) ** (1.0 / (alpha + beta))
        d = c / (self.compute_multiplier * n)
        return ComputeAllocation(c, n, d, self.loss(n, d))

    def tokens_for_target_loss(self, model_parameters: float, target_loss: float) -> float:
        """Return tokens needed for ``target_loss`` at fixed model size.

        The fitted law has a model-size floor
        ``E + A*N**(-alpha)``.  A target at or below that floor is unattainable for
        finite data and is rejected explicitly rather than producing a complex value.
        """

        n = _positive_finite("model_parameters", model_parameters)
        target = _positive_finite("target_loss", target_loss)
        residual = (
            target
            - self.irreducible_loss
            - self.model_coefficient * n ** (-self.model_exponent)
        )
        if residual <= 0.0:
            raise ValueError("target_loss is below the fitted floor for this model size")
        return (self.data_coefficient / residual) ** (1.0 / self.data_exponent)


@dataclass(frozen=True)
class BreakEvenResult:
    """Training and inference arithmetic for two equal-quality model choices."""

    target_loss: float
    small_training_flops: float
    large_training_tokens: float
    large_training_flops: float
    training_flops_saved_by_large_model: float
    extra_inference_flops_per_token_for_large_model: float
    break_even_generated_tokens: float


def inference_break_even(
    law: ScalingLaw,
    *,
    small_model_parameters: float,
    small_training_tokens: float,
    large_model_parameters: float,
    inference_flop_multiplier: float = 2.0,
) -> BreakEvenResult:
    """Compare an overtrained small model with an equal-quality larger model.

    The function assumes an arithmetic inference cost of
    ``inference_flop_multiplier * N`` per generated token.  It requires the larger
    model to use less training compute and more inference compute; otherwise a positive
    break-even does not exist under the declared scenario.
    """

    n_small = _positive_finite("small_model_parameters", small_model_parameters)
    d_small = _positive_finite("small_training_tokens", small_training_tokens)
    n_large = _positive_finite("large_model_parameters", large_model_parameters)
    if n_large <= n_small:
        raise ValueError("large_model_parameters must exceed small_model_parameters")
    infer_k = _positive_finite("inference_flop_multiplier", inference_flop_multiplier)

    target = law.loss(n_small, d_small)
    d_large = law.tokens_for_target_loss(n_large, target)
    c_small = law.training_flops(n_small, d_small)
    c_large = law.training_flops(n_large, d_large)
    training_saved = c_small - c_large
    inference_delta = infer_k * (n_large - n_small)
    if training_saved <= 0.0:
        raise ValueError("the larger model does not save training FLOPs in this scenario")
    return BreakEvenResult(
        target_loss=target,
        small_training_flops=c_small,
        large_training_tokens=d_large,
        large_training_flops=c_large,
        training_flops_saved_by_large_model=training_saved,
        extra_inference_flops_per_token_for_large_model=inference_delta,
        break_even_generated_tokens=training_saved / inference_delta,
    )


@dataclass(frozen=True)
class TrainingRunEstimate:
    """A transparent arithmetic runtime estimate for one pretraining run."""

    model_parameters: float
    training_tokens: float
    accelerator_count: int
    peak_flops_per_accelerator: float
    model_flop_utilization: float
    total_training_flops: float
    effective_cluster_flops_per_second: float
    estimated_seconds: float
    estimated_days: float
    training_tokens_per_second: float


def estimate_training_run(
    law: ScalingLaw,
    *,
    model_parameters: float,
    training_tokens: float,
    accelerator_count: int,
    peak_flops_per_accelerator: float,
    model_flop_utilization: float,
) -> TrainingRunEstimate:
    """Return an arithmetic duration estimate under a declared MFU assumption.

    Peak FLOPs are decimal FLOP/s.  The estimate excludes downtime, checkpoint stalls,
    data-loader starvation, communication imbalance, failures, and warm-up unless their
    effect is already folded into the supplied MFU.
    """

    n = _positive_finite("model_parameters", model_parameters)
    d = _positive_finite("training_tokens", training_tokens)
    devices = _positive_int("accelerator_count", accelerator_count)
    peak = _positive_finite("peak_flops_per_accelerator", peak_flops_per_accelerator)
    mfu = _fraction("model_flop_utilization", model_flop_utilization, allow_zero=False)
    total = law.training_flops(n, d)
    effective_rate = devices * peak * mfu
    seconds = total / effective_rate
    return TrainingRunEstimate(
        model_parameters=n,
        training_tokens=d,
        accelerator_count=devices,
        peak_flops_per_accelerator=peak,
        model_flop_utilization=mfu,
        total_training_flops=total,
        effective_cluster_flops_per_second=effective_rate,
        estimated_seconds=seconds,
        estimated_days=seconds / 86_400.0,
        training_tokens_per_second=d / seconds,
    )


law = ScalingLaw(**REFERENCE["scaling_law"])
for expected in REFERENCE["optimal_allocations"]:
    point = law.optimal_split(expected["compute_flops"])
    np.testing.assert_allclose(point.model_parameters, expected["model_parameters"], rtol=1e-11)
    np.testing.assert_allclose(point.training_tokens, expected["training_tokens"], rtol=1e-11)
    np.testing.assert_allclose(point.predicted_loss, expected["predicted_loss"], rtol=1e-11)

    # The derivative vanishes when the two marginal loss contributions balance.
    lhs = law.model_exponent * law.model_coefficient * point.model_parameters ** (-law.model_exponent)
    rhs = law.data_exponent * law.data_coefficient * point.training_tokens ** (-law.data_exponent)
    np.testing.assert_allclose(lhs, rhs, rtol=1e-11)

    # A local isoFLOP grid independently checks the analytic stationary point.
    grid_n = point.model_parameters * np.logspace(-0.4, 0.4, 401)
    grid_d = point.compute_flops / (law.compute_multiplier * grid_n)
    grid_loss = np.array([law.loss(n, d) for n, d in zip(grid_n, grid_d)])
    assert abs(int(np.argmin(grid_loss)) - 200) <= 1
    print(
        f"C={point.compute_flops:.2e}: N*={point.model_parameters/1e9:7.2f}B, "
        f"D*={point.training_tokens/1e12:7.3f}T, D/N={point.tokens_per_parameter:7.2f}"
    )

small_ref = REFERENCE["small_model_scenario"]
small_loss = law.loss(small_ref["model_parameters"], small_ref["training_tokens"])
same_budget = law.optimal_split(small_ref["training_flops"])
np.testing.assert_allclose(small_loss, small_ref["predicted_loss"], rtol=1e-11)
np.testing.assert_allclose(
    small_loss - same_budget.predicted_loss,
    small_ref["loss_gap_to_same_budget_optimum"],
    rtol=1e-11,
)
print(
    f"8.03B/15T: D/N={small_ref['tokens_per_parameter']:.1f}, "
    f"same-budget loss gap={small_ref['loss_gap_to_same_budget_optimum']:.6f}"
)
print('Scaling-law optimum, isoFLOP grid, and first-order condition verified')

small_ref = REFERENCE["small_model_scenario"]
break_ref = REFERENCE["break_even"]
break_even = inference_break_even(
    law,
    small_model_parameters=small_ref["model_parameters"],
    small_training_tokens=small_ref["training_tokens"],
    large_model_parameters=break_ref["large_model_parameters"],
    inference_flop_multiplier=2.0,
)
for field in (
    "target_loss",
    "small_training_flops",
    "large_training_tokens",
    "large_training_flops",
    "training_flops_saved_by_large_model",
    "extra_inference_flops_per_token_for_large_model",
    "break_even_generated_tokens",
):
    np.testing.assert_allclose(getattr(break_even, field), break_ref[field], rtol=1e-11)

# At T*, both equal-quality plans have the same declared lifecycle arithmetic.
t_star = break_even.break_even_generated_tokens
small_total = break_even.small_training_flops + 2.0 * small_ref["model_parameters"] * t_star
large_total = break_even.large_training_flops + 2.0 * break_ref["large_model_parameters"] * t_star
np.testing.assert_allclose(small_total, large_total, rtol=1e-12)

run_ref = REFERENCE["training_run"]
run = estimate_training_run(
    law,
    model_parameters=run_ref["model_parameters"],
    training_tokens=run_ref["training_tokens"],
    accelerator_count=run_ref["accelerator_count"],
    peak_flops_per_accelerator=run_ref["peak_flops_per_accelerator"],
    model_flop_utilization=run_ref["model_flop_utilization"],
)
for field in (
    "total_training_flops",
    "effective_cluster_flops_per_second",
    "estimated_seconds",
    "estimated_days",
    "training_tokens_per_second",
):
    np.testing.assert_allclose(getattr(run, field), run_ref[field], rtol=1e-12)

print(f"equal-loss 70B tokens: {break_even.large_training_tokens/1e12:.3f}T")
print(f"break-even: {t_star/1e12:.3f}T generated tokens")
print(
    f"arithmetic run scenario: {run.accelerator_count} × "
    f"{run.peak_flops_per_accelerator/1e15:.1f} PFLOP/s peak × "
    f"MFU {run.model_flop_utilization:.0%} → {run.estimated_days:.3f} days"
)
print('Equal-quality break-even and training-time scenario verified')

@dataclass(frozen=True)
class DataSlice:
    """One unique-token pool and its unnormalized sampling weight."""

    name: str
    unique_tokens: float
    sampling_weight: float

    def __post_init__(self) -> None:
        if not isinstance(self.name, str) or not self.name.strip():
            raise ValueError("name must be a non-empty string")
        _positive_finite("unique_tokens", self.unique_tokens)
        _positive_finite("sampling_weight", self.sampling_weight)


@dataclass(frozen=True)
class MixtureAllocation:
    """Allocated tokens and effective epochs for one data slice."""

    name: str
    normalized_weight: float
    allocated_tokens: float
    unique_tokens: float
    effective_epochs: float


def allocate_data_mixture(
    slices: Sequence[DataSlice], total_training_tokens: float
) -> tuple[MixtureAllocation, ...]:
    """Allocate a token budget by normalized weights and report effective epochs."""

    total = _positive_finite("total_training_tokens", total_training_tokens)
    items = tuple(slices)
    if not items:
        raise ValueError("slices must be non-empty")
    if not all(isinstance(item, DataSlice) for item in items):
        raise TypeError("every slice must be a DataSlice")
    names = [item.name for item in items]
    if len(names) != len(set(names)):
        raise ValueError("slice names must be unique")
    weight_sum = math.fsum(item.sampling_weight for item in items)
    allocations = tuple(
        MixtureAllocation(
            name=item.name,
            normalized_weight=item.sampling_weight / weight_sum,
            allocated_tokens=total * item.sampling_weight / weight_sum,
            unique_tokens=item.unique_tokens,
            effective_epochs=(total * item.sampling_weight / weight_sum) / item.unique_tokens,
        )
        for item in items
    )
    if not math.isclose(
        math.fsum(item.allocated_tokens for item in allocations),
        total,
        rel_tol=1e-12,
        abs_tol=1e-6,
    ):
        raise AssertionError("mixture allocations do not sum to the requested token budget")
    return allocations


mix_ref = REFERENCE["data_mixture"]
slices = tuple(DataSlice(**item) for item in mix_ref["slices"])
allocations = allocate_data_mixture(slices, mix_ref["total_training_tokens"])
assert len(allocations) == len(mix_ref["allocations"])
for observed, expected in zip(allocations, mix_ref["allocations"]):
    assert observed.name == expected["name"]
    for field in ("normalized_weight", "allocated_tokens", "unique_tokens", "effective_epochs"):
        np.testing.assert_allclose(getattr(observed, field), expected[field], rtol=1e-12)
assert math.isclose(sum(item.normalized_weight for item in allocations), 1.0, rel_tol=0.0, abs_tol=1e-12)
assert math.isclose(
    sum(item.allocated_tokens for item in allocations),
    mix_ref["total_training_tokens"],
    rel_tol=1e-12,
)

for item in allocations:
    print(
        f"{item.name:6s}: weight={item.normalized_weight:5.1%}, "
        f"allocated={item.allocated_tokens/1e12:5.2f}T, epochs={item.effective_epochs:5.3f}"
    )
print('Mixture allocation and effective epochs verified')

@dataclass(frozen=True)
class SpectralSummary:
    """Singular-value and orthogonality diagnostics for a matrix update."""

    minimum_singular_value: float
    maximum_singular_value: float
    condition_number: float
    orthogonality_error_frobenius: float


def spectral_summary(matrix: Sequence[Sequence[float]]) -> SpectralSummary:
    """Return diagnostics against the identity on the smaller matrix dimension."""

    array = np.asarray(matrix, dtype=np.float64)
    if array.ndim != 2 or 0 in array.shape:
        raise ValueError("matrix must be a non-empty two-dimensional array")
    if not np.isfinite(array).all():
        raise ValueError("matrix must contain only finite values")
    singular = np.linalg.svd(array, compute_uv=False)
    minimum = float(singular.min())
    maximum = float(singular.max())
    condition = math.inf if minimum == 0.0 else maximum / minimum
    if array.shape[0] >= array.shape[1]:
        gram = array.T @ array
        identity = np.eye(array.shape[1], dtype=np.float64)
    else:
        gram = array @ array.T
        identity = np.eye(array.shape[0], dtype=np.float64)
    error = float(np.linalg.norm(gram - identity, ord="fro"))
    return SpectralSummary(minimum, maximum, condition, error)


def newton_schulz_orthogonalize(
    matrix: Sequence[Sequence[float]],
    *,
    quintic_steps: int = 5,
    cubic_steps: int = 0,
) -> np.ndarray:
    """Return a rectangular Newton--Schulz polar-factor approximation.

    The matrix is Frobenius-normalized before the iterations.  Tall matrices are
    transposed so the left-Gram matrix uses the smaller dimension, then transposed back.  The five
    quintic steps reproduce the Muon teaching trace; optional cubic steps illustrate
    how a costlier classical tail tightens the singular-value band.
    """

    array = np.asarray(matrix, dtype=np.float64)
    if array.ndim != 2 or 0 in array.shape:
        raise ValueError("matrix must be a non-empty two-dimensional array")
    if not np.isfinite(array).all():
        raise ValueError("matrix must contain only finite values")
    q_steps = _non_negative_int("quintic_steps", quintic_steps)
    c_steps = _non_negative_int("cubic_steps", cubic_steps)
    scale = float(np.max(np.abs(array)))
    if scale == 0.0:
        raise ValueError("matrix must not be all zeros")

    transposed = array.shape[0] > array.shape[1]
    work = array.T.copy() if transposed else array.copy()
    work /= scale
    work /= np.linalg.norm(work, ord="fro")
    for _ in range(q_steps):
        gram = work @ work.T
        work = 3.4445 * work - 4.7750 * gram @ work + 2.0315 * gram @ gram @ work
    for _ in range(c_steps):
        work = 1.5 * work - 0.5 * (work @ work.T) @ work
    return work.T if transposed else work


@dataclass(frozen=True)
class QKClipResult:
    """Per-head QK-Clip scales for one monitoring step."""

    max_logits: np.ndarray
    threshold: float
    split: float
    gamma: np.ndarray
    query_scale: np.ndarray
    key_scale: np.ndarray
    corrected_max_logits: np.ndarray


def qk_clip_scales(
    max_logits: Sequence[float],
    threshold: float,
    *,
    split: float = 0.5,
) -> QKClipResult:
    """Return per-head parameter scales that restore monitored maxima to ``threshold``.

    ``split`` assigns ``gamma**split`` to Q and ``gamma**(1-split)`` to K.  The
    correction is an intervention at the current step; continued training can regrow
    the weights, so monitoring and reapplication remain necessary.
    """

    values = np.asarray(max_logits, dtype=np.float64)
    if values.ndim != 1 or values.size == 0:
        raise ValueError("max_logits must be a non-empty one-dimensional vector")
    if not np.isfinite(values).all() or np.any(values < 0.0):
        raise ValueError("max_logits must contain finite non-negative values")
    tau = _positive_finite("threshold", threshold)
    share = _fraction("split", split)
    gamma = np.ones_like(values)
    over = values > tau
    gamma[over] = tau / values[over]
    q_scale = gamma**share
    k_scale = gamma ** (1.0 - share)
    corrected = values * q_scale * k_scale
    if np.any(corrected > tau * (1.0 + 1e-12)):
        raise AssertionError("QK-Clip failed to satisfy the threshold")
    return QKClipResult(values.copy(), tau, share, gamma, q_scale, k_scale, corrected)


mu_ref = REFERENCE["muon"]
rng = np.random.default_rng(mu_ref["seed"])
gradient = rng.standard_normal(tuple(mu_ref["shape"]))
before = spectral_summary(gradient)
quintic = newton_schulz_orthogonalize(
    gradient,
    quintic_steps=mu_ref["quintic_steps"],
    cubic_steps=0,
)
after_quintic = spectral_summary(quintic)
with_cubic = newton_schulz_orthogonalize(
    gradient,
    quintic_steps=mu_ref["quintic_steps"],
    cubic_steps=mu_ref["cubic_steps"],
)
after_cubic = spectral_summary(with_cubic)

# SVD last bits can vary across BLAS builds; mechanism-level checks use 1e-6 tolerance.
for observed, key in (
    (before, "before"),
    (after_quintic, "after_quintic"),
    (after_cubic, "after_cubic_tail"),
):
    expected = mu_ref[key]
    for field in (
        "minimum_singular_value",
        "maximum_singular_value",
        "condition_number",
        "orthogonality_error_frobenius",
    ):
        np.testing.assert_allclose(getattr(observed, field), expected[field], rtol=1e-6, atol=1e-9)

wide_result = newton_schulz_orthogonalize(
    gradient.T,
    quintic_steps=mu_ref["quintic_steps"],
    cubic_steps=0,
)
assert list(wide_result.shape) == mu_ref["wide_shape"]
assert np.allclose(wide_result, quintic.T, atol=1e-12, rtol=0.0)
assert after_cubic.orthogonality_error_frobenius < after_quintic.orthogonality_error_frobenius

print(f"initial condition ratio: {before.condition_number:.3f}")
print(
    f"quintic band: [{after_quintic.minimum_singular_value:.6f}, "
    f"{after_quintic.maximum_singular_value:.6f}], "
    f"error={after_quintic.orthogonality_error_frobenius:.6f}"
)
print(
    f"+ cubic tail: [{after_cubic.minimum_singular_value:.6f}, "
    f"{after_cubic.maximum_singular_value:.6f}], "
    f"error={after_cubic.orthogonality_error_frobenius:.6f}"
)
print('Newton–Schulz shape/band/error contracts verified')

qk_ref = REFERENCE["qk_clip"]
qk = qk_clip_scales(qk_ref["max_logits"], qk_ref["threshold"], split=qk_ref["split"])
for observed, key in (
    (qk.gamma, "gamma"),
    (qk.query_scale, "query_scale"),
    (qk.key_scale, "key_scale"),
    (qk.corrected_max_logits, "corrected_max_logits"),
):
    np.testing.assert_allclose(observed, qk_ref[key], rtol=1e-12, atol=1e-12)
np.testing.assert_allclose(qk.query_scale * qk.key_scale, qk.gamma, rtol=1e-12)
assert np.all(qk.corrected_max_logits <= qk.threshold * (1.0 + 1e-12))

# A clipped display hides the excess while the source values continue to grow.
regrown = qk.corrected_max_logits * 1.1**3
clipped_view = np.minimum(regrown, qk.threshold)
np.testing.assert_allclose(regrown, qk_ref["regrown_after_three_10pct_steps"], rtol=1e-12)
np.testing.assert_allclose(clipped_view, qk_ref["clipped_view_after_regrowth"], rtol=1e-12)
assert regrown.max() > qk.corrected_max_logits.max()

# Reapplying the parameter correction after regrowth restores the monitored threshold.
reapplied = qk_clip_scales(regrown, qk.threshold, split=qk.split)
np.testing.assert_allclose(
    reapplied.corrected_max_logits,
    qk_ref["reapplied_corrected_max_logits"],
    rtol=1e-12,
    atol=1e-12,
)

print("gamma:", qk.gamma)
print("Q scales:", qk.query_scale)
print("K scales:", qk.key_scale)
print("corrected maxima:", qk.corrected_max_logits)
print("regrown source after three +10% steps:", regrown)
print('QK-Clip correction and regrowth trace verified')

@dataclass(frozen=True)
class AdamWStatePolicy:
    """Per-parameter byte components for mixed-precision AdamW state accounting."""

    model_weight_bytes: int = 2
    gradient_bytes: int = 2
    master_weight_bytes: int = 4
    first_moment_bytes: int = 4
    second_moment_bytes: int = 4

    def __post_init__(self) -> None:
        for name, value in (
            ("model_weight_bytes", self.model_weight_bytes),
            ("gradient_bytes", self.gradient_bytes),
            ("master_weight_bytes", self.master_weight_bytes),
            ("first_moment_bytes", self.first_moment_bytes),
            ("second_moment_bytes", self.second_moment_bytes),
        ):
            _positive_int(name, value)

    @property
    def optimizer_state_bytes(self) -> int:
        """Return master weights plus first and second moments per parameter."""

        return self.master_weight_bytes + self.first_moment_bytes + self.second_moment_bytes

    @property
    def total_bytes_per_parameter(self) -> int:
        """Return all declared state bytes per parameter."""

        return self.model_weight_bytes + self.gradient_bytes + self.optimizer_state_bytes


def zero_state_bytes_per_rank(
    parameter_count: int,
    world_size: int,
    zero_stage: int,
    policy: AdamWStatePolicy | None = None,
) -> float:
    """Return the idealized parameter-state bytes per rank for ZeRO stages 0--3.

    Stage 0 replicates everything. Stage 1 shards optimizer state, stage 2 also shards
    gradients, and stage 3 also shards model weights. Activations, communication buffers,
    temporary workspaces, fragmentation, and framework metadata are deliberately absent.
    """

    parameters = _positive_int("parameter_count", parameter_count)
    world = _positive_int("world_size", world_size)
    stage = _non_negative_int("zero_stage", zero_stage)
    if stage > 3:
        raise ValueError("zero_stage must be one of 0, 1, 2, 3")
    p = policy or AdamWStatePolicy()
    weights = parameters * p.model_weight_bytes
    gradients = parameters * p.gradient_bytes
    optimizer = parameters * p.optimizer_state_bytes
    if stage == 0:
        return float(weights + gradients + optimizer)
    if stage == 1:
        return float(weights + gradients + optimizer / world)
    if stage == 2:
        return float(weights + (gradients + optimizer) / world)
    return float((weights + gradients + optimizer) / world)


def pipeline_bubble_fraction(stages: int, microbatches: int) -> float:
    """Return the idealized 1F1B bubble fraction ``(p-1)/(m+p-1)``."""

    p = _positive_int("stages", stages)
    m = _positive_int("microbatches", microbatches)
    return (p - 1) / (m + p - 1)


def minimum_microbatches_for_bubble(stages: int, maximum_bubble: float) -> int:
    """Return the smallest integer ``m`` whose idealized bubble is below the limit."""

    p = _positive_int("stages", stages)
    limit = _fraction("maximum_bubble", maximum_bubble, allow_zero=False)
    if p == 1:
        return 1
    candidate = max(1, math.ceil((p - 1) * (1.0 / limit - 1.0)))
    while pipeline_bubble_fraction(p, candidate) >= limit:
        candidate += 1
    while candidate > 1 and pipeline_bubble_fraction(p, candidate - 1) < limit:
        candidate -= 1
    return candidate


memory_ref = REFERENCE["training_state_memory"]
policy = AdamWStatePolicy(
    model_weight_bytes=memory_ref["policy"]["model_weight_bytes"],
    gradient_bytes=memory_ref["policy"]["gradient_bytes"],
    master_weight_bytes=memory_ref["policy"]["master_weight_bytes"],
    first_moment_bytes=memory_ref["policy"]["first_moment_bytes"],
    second_moment_bytes=memory_ref["policy"]["second_moment_bytes"],
)
assert policy.total_bytes_per_parameter == memory_ref["policy"]["total_bytes_per_parameter"] == 16
for expected in memory_ref["rows"]:
    byte_count = zero_state_bytes_per_rank(
        memory_ref["parameter_count"],
        memory_ref["world_size"],
        expected["zero_stage"],
        policy,
    )
    np.testing.assert_allclose(byte_count, expected["bytes_per_rank"], rtol=0.0, atol=0.0)
    np.testing.assert_allclose(bytes_to_decimal_gb(byte_count), expected["decimal_gb_per_rank"], rtol=1e-12)
    np.testing.assert_allclose(bytes_to_gib(byte_count), expected["gib_per_rank"], rtol=1e-12)
    print(
        f"ZeRO-{expected['zero_stage']}: {byte_count/1e9:8.4f} GB/rank = "
        f"{byte_count/2**30:8.4f} GiB/rank"
    )

assert memory_ref["rows"][0]["decimal_gb_per_rank"] == 112.0
assert memory_ref["rows"][3]["decimal_gb_per_rank"] == 1.75
print('ZeRO state payload in GB/GiB verified')

def expect_raises(error_type: type[BaseException], fn, *args, **kwargs) -> None:
    """Assert that a negative test raises the declared exception type."""

    try:
        fn(*args, **kwargs)
    except error_type:
        return
    raise AssertionError(f"expected {error_type.__name__} from {getattr(fn, '__name__', fn)!r}")


pipeline_ref = REFERENCE["pipeline"]
for expected in pipeline_ref["rows"]:
    fraction = pipeline_bubble_fraction(expected["stages"], expected["microbatches"])
    np.testing.assert_allclose(fraction, expected["bubble_fraction"], rtol=0.0, atol=1e-15)
    print(
        f"p={expected['stages']}, m={expected['microbatches']:3d}: "
        f"bubble={fraction:.3%}, ideal efficiency={1-fraction:.3%}"
    )
minimum = minimum_microbatches_for_bubble(8, 0.05)
assert minimum == pipeline_ref["minimum_microbatches_for_5pct"] == 134
assert pipeline_bubble_fraction(8, minimum) < 0.05
assert pipeline_bubble_fraction(8, minimum - 1) >= 0.05
print(f"minimum m for bubble <5% at p=8: {minimum}")

# Negative tests make failure modes part of the learner-visible contract.
expect_raises(ValueError, law.optimal_split, 0.0)
expect_raises(ValueError, law.loss, -1.0, 1.0e9)
expect_raises(ValueError, law.tokens_for_target_loss, 1.0e6, law.irreducible_loss)
expect_raises(ValueError, estimate_training_run, law, model_parameters=1e9, training_tokens=1e9,
              accelerator_count=8, peak_flops_per_accelerator=1e14, model_flop_utilization=1.1)
expect_raises(ValueError, allocate_data_mixture, tuple(), 1.0e9)
expect_raises(ValueError, allocate_data_mixture,
              (DataSlice("dup", 1e9, 1.0), DataSlice("dup", 1e9, 1.0)), 1e9)
expect_raises(ValueError, newton_schulz_orthogonalize, np.zeros((2, 2)))
expect_raises(ValueError, qk_clip_scales, [-1.0], 100.0)
expect_raises(ValueError, qk_clip_scales, [1.0], 100.0, split=1.1)
expect_raises(ValueError, zero_state_bytes_per_rank, 1_000, 8, 4)
expect_raises(ValueError, pipeline_bubble_fraction, 0, 32)
expect_raises(ValueError, minimum_microbatches_for_bubble, 8, 0.0)

print('Pipeline bubble and negative tests passed')

"""Two explicit AdamW steps, with supplied gradients and no training framework."""
import numpy as np


def adamw_trace():
    theta = np.array([1., -2.])
    m, v = np.zeros(2), np.zeros(2)
    eta, decay, beta1, beta2, eps = .01, .1, .9, .999, 1e-8
    history = [theta.copy()]
    for step, gradient in enumerate((np.array([.2, -.4]), np.array([-.1, .2])), 1):
        m = beta1 * m + (1 - beta1) * gradient
        v = beta2 * v + (1 - beta2) * gradient**2
        m_hat, v_hat = m / (1 - beta1**step), v / (1 - beta2**step)
        theta = (1 - eta * decay) * theta - eta * m_hat / (np.sqrt(v_hat) + eps)
        history.append(theta.copy())
    return np.array(history)


trace = adamw_trace()
np.testing.assert_allclose(trace[1], [.9890000005, -1.98800000025], atol=1e-12, rtol=0)
# The second gradient reverses, but accumulated momentum retains its old sign.
assert trace[2, 0] < trace[1, 0] and trace[2, 1] > trace[1, 1]
print("AdamW parameter trace:", trace)

def check_minimal_core(reference_qk_clip_scales):
    """Minimal mechanisms embedded verbatim in the lecture and notebook."""
    from dataclasses import dataclass
    import math
    from numbers import Real
    import numpy as np


    def positive(name, value):
        if isinstance(value, (bool, np.bool_)) or not isinstance(value, Real):
            raise ValueError(f"{name} must be a real number")
        value = float(value)
        if not math.isfinite(value) or value <= 0:
            raise ValueError(f"{name} must be positive and finite")
        return value


    @dataclass(frozen=True)
    class ScalingLaw:
        E: float = 1.69
        A: float = 406.4
        B: float = 410.7
        alpha: float = 0.34
        beta: float = 0.28
        kappa: float = 6.0

        def __post_init__(self):
            for name in ("A", "B", "alpha", "beta", "kappa"):
                object.__setattr__(self, name, positive(name, getattr(self, name)))
            if isinstance(self.E, (bool, np.bool_)) or not isinstance(self.E, Real):
                raise ValueError("E must be a real number")
            if not math.isfinite(self.E) or self.E < 0:
                raise ValueError("E must be finite and non-negative")
            object.__setattr__(self, "E", float(self.E))

        def loss(self, N, D):
            N, D = positive("N", N), positive("D", D)
            return self.E + self.A * N**(-self.alpha) + self.B * D**(-self.beta)

        def optimal_split(self, C):
            C = positive("C", C)
            N = (
                (self.alpha * self.A) / (self.beta * self.B)
                * (C / self.kappa) ** self.beta
            ) ** (1 / (self.alpha + self.beta))
            return N, C / (self.kappa * N)


    def newton_schulz(G, quintic_steps=5):
        """Trace a rectangular polar-factor approximation."""
        X = np.asarray(G, dtype=np.float64)
        if X.ndim != 2 or 0 in X.shape or not np.isfinite(X).all():
            raise ValueError("G must be a finite non-empty matrix")
        if isinstance(quintic_steps, bool) or not isinstance(quintic_steps, int) or quintic_steps < 0:
            raise ValueError("quintic_steps must be a non-negative integer")
        scale = np.max(np.abs(X))
        if scale == 0:
            raise ValueError("G must not be all zeros")
        transposed = X.shape[0] > X.shape[1]
        X = (X.T if transposed else X) / scale
        X /= np.linalg.norm(X, ord="fro")
        for _ in range(quintic_steps):
            XXt = X @ X.T
            X = 3.4445 * X - 4.7750 * XXt @ X + 2.0315 * XXt @ XXt @ X
        return X.T if transposed else X


    def qk_clip_scales(max_logits, tau=100.0, split=0.5):
        """Return gamma and Q/K scales for one monitoring step."""
        maxima = np.asarray(max_logits, dtype=np.float64)
        if maxima.ndim != 1 or maxima.size == 0 or not np.isfinite(maxima).all():
            raise ValueError("max_logits must be a finite vector")
        if np.any(maxima < 0) or not 0 <= split <= 1:
            raise ValueError("maxima must be non-negative and split must lie in [0, 1]")
        tau = positive("tau", tau)
        gamma = np.ones_like(maxima)
        over = maxima > tau
        gamma[over] = tau / maxima[over]
        q_scale = gamma**split
        k_scale = gamma**(1 - split)
        return gamma, q_scale, k_scale, maxima * q_scale * k_scale


    # Minimal-core coefficients must be real scalars, never parseable strings.
    for field in ("E", "A", "B", "alpha", "beta", "kappa"):
        for bad in ("1.0", True, np.bool_(True), complex(1, 0), np.nan, np.inf):
            expect_raises(ValueError, ScalingLaw, **{field: bad})
    assert isinstance(ScalingLaw(A=np.float32(406.4)).A, float)
    assert ScalingLaw(E=0).loss(1e9, 1e12) > 0
    np.testing.assert_allclose(newton_schulz(gradient), quintic, atol=1e-12)
    for factor in (1e-300, 1e300):
        np.testing.assert_allclose(newton_schulz_orthogonalize(gradient * factor), quintic, atol=1e-12)
    # A scalar polynomial applied to SVD singular values is an independent oracle.
    u, singular, vt = np.linalg.svd(gradient, full_matrices=False)
    z = singular / np.linalg.norm(singular)
    for _ in range(5):
        z = 3.4445*z - 4.7750*z**3 + 2.0315*z**5
    np.testing.assert_allclose(quintic, (u*z)@vt, atol=2e-12)
    rank_one = np.array([[1., 0.], [0., 0.], [0., 0.]])
    assert np.linalg.matrix_rank(newton_schulz_orthogonalize(rank_one)) == 1
    # Test the bilinear mechanism on actual Q/K arrays, not just multiplied maxima.
    rng = np.random.default_rng(17)
    q, k = rng.normal(size=(4, 5, 8))*10, rng.normal(size=(4, 7, 8))*10
    logits = np.einsum("htd,hsd->hts", q, k) / np.sqrt(8)
    maxima = np.maximum(logits.max(axis=(1,2)), 0)
    scales = reference_qk_clip_scales(maxima, 100, split=.3)
    corrected = np.einsum("htd,hsd->hts", q*scales.query_scale[:,None,None],
                          k*scales.key_scale[:,None,None])/np.sqrt(8)
    np.testing.assert_allclose(corrected, logits*scales.gamma[:,None,None])
    assert corrected.max() <= 100 + 1e-10
    zero_clip = reference_qk_clip_scales([0., 100., 200.], 100.)
    np.testing.assert_allclose(zero_clip.gamma, [1., 1., .5])
    np.testing.assert_allclose(qk_clip_scales([0., 100., 200.], 100.)[0], zero_clip.gamma)
    assert pipeline_bubble_fraction(8, 133) == .05
    assert minimum_microbatches_for_bubble(8, .05) == 134
    print("Independent mechanism checks: PASS")


check_minimal_core(qk_clip_scales)

from __future__ import annotations
import importlib.util, subprocess, sys
for package in ("numpy", "matplotlib"):
    if importlib.util.find_spec(package) is None:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", package])


import math
import random
import time
from collections import Counter, defaultdict

import matplotlib.pyplot as plt
import numpy as np

SEED = 20260801
random.seed(SEED)
np.random.seed(SEED)
print(f"reproducible smoke tier; seed={SEED}")
DEMO_LANGUAGE = 'en'

plt.rcParams.update({"font.family":"DejaVu Sans", "font.size":12, "axes.titlesize":15, "axes.labelsize":13, "xtick.labelsize":11, "ytick.labelsize":11, "figure.dpi":150})

import re
from collections import Counter
# Synthetic documents; language and quality labels are supplied, not inferred.
docs=[
 ("a","en",.9,"matrix updates use gradients to change model weights"),
 ("b","en",.8,"MATRIX updates use gradients to change model weights"),
 ("c","en",.85,"matrix updates use gradients to change neural weights"),
 ("d","en",.1,"buy buy buy buy buy buy buy buy"),
 ("e","ru",.9,"Градиенты помогают обновлять параметры модели при обучении"),
 ("f","xx",.8,"unknown script placeholder text for this example"),
 ("g","en",.8,"Email example.person@example.test for the synthetic exercise"),
 ("h","en",.9,"attention reads values by comparing queries and keys"),
]
def words(text): return re.findall(r"\w+",text.casefold())
def shingles(text,k=2):
 w=words(text);return {tuple(w[i:i+k]) for i in range(len(w)-k+1)}
def jaccard(a,b): return len(a&b)/len(a|b) if a|b else 1.
def curate(order):
 kept=[];counts=Counter();seen=[]
 for docid,lang,quality,text in docs:
  failed=None
  for stage in order:
   if stage=="language" and lang not in {"en","ru"}:failed=stage
   elif stage=="quality" and quality<.5:failed=stage
   elif stage=="pii" and re.search(r"\S+@\S+",text):failed=stage
   elif stage=="dedup" and any(jaccard(shingles(text),s)>=.5 for s in seen):failed=stage
   if failed:break
  if failed:counts[failed]+=1
  else:kept.append(docid);seen.append(shingles(text))
 counts["kept"]=len(kept)
 assert sum(counts.values())==len(docs)
 return kept,counts
kept,counts=curate(["language","quality","pii","dedup"])
assert kept==["a","e","h"]
print("survival:",counts,"retained:",kept)
print("near-duplicate Jaccard:",jaccard(shingles(docs[0][3]),shingles(docs[2][3])))
# Cross-language paraphrases are not detected by this lexical mechanism.
assert jaccard(shingles(docs[0][3]),shingles(docs[4][3])) == 0
fig,ax=plt.subplots(figsize=(11,4.5),layout="constrained")
ax.bar(list(counts),list(counts.values()),color="#5E81AC")
ax.set_ylabel("Документы" if DEMO_LANGUAGE=="ru" else "Documents")
ax.set_title("Исключающие причины отсева" if DEMO_LANGUAGE=="ru" else "Exclusive retention/rejection counts")
plt.show()

# Simulate two ranks on CPU: all-gather weights, reduce-scatter gradients,
# update owned shards. No actual network, FSDP runtime, or GPU memory benchmark.
rng=np.random.default_rng(SEED)
x=rng.normal(size=(12,8));y=rng.normal(size=(12,4));w=rng.normal(size=(8,4))
shards=[a.copy() for a in np.array_split(w,2,axis=0)]
lr=.03
for step in range(4):
 full=np.concatenate(shards,axis=0)
 # Equal local batch sizes; average local gradients gives the global gradient.
 local=[]
 for xb,yb in zip(np.array_split(x,2),np.array_split(y,2)):
  local.append(xb.T@(xb@full-yb)/len(xb))
 grad=sum(local)/2
 refgrad=x.T@(x@w-y)/len(x)
 np.testing.assert_allclose(grad,refgrad,atol=1e-12)
 for i,g in enumerate(np.array_split(grad,2,axis=0)):shards[i]-=lr*g
 w-=lr*refgrad
 np.testing.assert_allclose(np.concatenate(shards),w,atol=1e-12)
assert all(s.nbytes==w.nbytes//2 for s in shards)
print("CPU shard parity: PASS; owned weight bytes/rank:",shards[0].nbytes,
      "temporary gathered weights:",full.nbytes)
# Different ledger: hypothetical mixed-precision AdamW, not this NumPy dtype.
P=1_000_000_000
for world in (1,2,8):
 for stage in (0,1,2,3):
  # Per parameter: weights 2 B, gradients 2 B, master+moments 12 B.
  per_parameter=(16,4+12/world,2+14/world,16/world)[stage]
  payload=P*per_parameter
  print(f"ZeRO-{stage}, {world} ranks: {payload/1e9:.3f} GB/rank")
 # An illustrative additional gather window and live activation reserve.
 gather=100_000_000*2 if world>1 else 0
 activations=500_000_000
 with_buffers=P*16/world+gather+activations
 print(f"ZeRO-3 plus declared buffers: {with_buffers/1e9:.3f} GB/rank")
print("Activations, gathers, communication buffers, and allocator reserve are additional.")