src.tests.spconv3d_mixed_precision_e2e_test

Classes

SpConv3DMixedPrecisionE2ETest()

class src.tests.spconv3d_mixed_precision_e2e_test.SpConv3DMixedPrecisionE2ETest
Author:

Alberto M. Esmoris Pena

End-to-end mixed-precision audit for the SpConv stack — complements the in-isolation fp16 / bf16 coverage of MaskedBatchNormalizationTest by exercising the full pipeline (preprocessor + sequencer + arch + loss) under Keras 3 mixed_float16 and mixed_bfloat16 global policies. This test class validates the full pipeline under mixed_float16 and mixed_bfloat16 policies on a tiny synthetic batch: build the model, save/restore the policy, forward pass, gradient check, loss curve sanity.

Foundation-scale training depends on the framework being fp16/bf16-safe end-to-end. This gate catches regressions in:

  • Per-layer dtype handling (SubmanifoldSpConv3DLayer, DownsamplingSpConv3DLayer, UpsamplingSpConv3DLayer, SpConv3DEncodingLayer, SpConv3DDecodingLayer, DenseWithGroundRow).

  • Output head softmax stability (C.3 pins it to fp32 via Dense(dtype='float32')).

  • Gradient signal under fp16 (compute throughput halves; loss scaling kicks in implicitly via Keras’ policy machinery).

Subtests:

  1. subtest_fp32_baseline_fits_one_epoch — control.

  2. subtest_mixed_float16_fits_one_epoch.

  3. subtest_mixed_bfloat16_fits_one_epoch.

  4. subtest_mixed_float16_output_dtype_fp32 — softmax stays fp32.

  5. subtest_mixed_float16_gradient_signal_nonzero — every trainable kernel sees a non-zero gradient after one step.

  6. subtest_mixed_float16_forward_no_naninf.

  7. subtest_mixed_float16_loss_curve_matches_fp32 — final loss within 5x of fp32 baseline.

__init__()

Basic configuration for any VL3D test.

Parameters:

name (str) – Test name

run()

Run the test.

Returns:

True if test is successfully passed, False otherwise.

Return type:

bool

subtest_fp32_baseline_fits_one_epoch()

Control: smoke arch + tiny synthetic batch under fp32.

subtest_mixed_float16_fits_one_epoch()

Under mixed_float16, a one-epoch fit must finish with a finite loss.

subtest_mixed_bfloat16_fits_one_epoch()

Same contract under mixed_bfloat16.

subtest_mixed_float16_output_dtype_fp32()

The output head must stay fp32 under mixed_float16 so the softmax does not overflow. C.3 pins cwise_out via dtype='float32'.

subtest_mixed_float16_gradient_signal_nonzero()

After one step under mixed_float16, every trainable kernel must have a non-zero gradient norm (>1e-8). Underflow would zero the gradient and stall training at scale.

subtest_mixed_float16_forward_no_naninf()

Forward pass under mixed_float16 produces finite output on a synthetic batch (no NaN / Inf from variance overflow or fp16 underflow).

subtest_mixed_float16_loss_curve_matches_fp32()

3-epoch fit under mixed_float16 produces a finite final loss within 5x of the fp32 baseline. The 5x tolerance absorbs implicit loss-scaling differences without accepting outright divergence.