LamTNguyen/archive-TPUv5p-16-2026-08-26

Archive TPUv5p-16 — 2026-08-26

Public research archive for the LARA / Self-Flow experiments run on a Google Cloud TPU v5p-16 slice (two workers, eight TPU chips). The snapshot was created on 2026-08-26 (Asia/HoChiMinh) before the VM storage was cleaned.

This is a reproducibility archive, not a polished model release. It includes source snapshots, run scripts, raw JSON/CSV/log outputs, preview images, generated latent batches, training checkpoints, input latent datasets, evaluation references, and model weights used by the runs. See MANIFEST.sha256 and MANIFEST.tsv for file-level verification.

Archive layout

The Python virtual environments and package caches are intentionally not archived: they are rebuildable, very large, and not experiment data. Credential files and authentication tokens are also excluded.

Main experiments and results

Detailed reports live under code/Self-Flow-local/outputs/**/RESULTS.md and code/Self-Flow-worker1/outputs/**/RESULTS.md. The headline results below are copied from those reports.

E-DiT / Elastic-DiT on LARA-B (SiT-B)

The implementation adds timestep-conditioned routers, hard-forward/STE-backward block and MLP-width gates, compute-budget and margin losses, preplanned masks, residual caching, calibrated static route buckets, and a single route-level TPU switch without per-block lax.cond.

Training used 10,000 post-training steps, global batch 256, and all eight TPU cores. Final router statistics included a 0.5397 hard skip ratio and finite updates throughout. The W&B run is al5ui419.

The shared-bucket static TPU serving path measured 1.105 ms/image (904.9 images/s), a 1.83x denoiser-only speedup over the 2.023 ms/image full static control. This fast path uses one aggregate route bucket for a global batch of 256; sample-wise routing was used for quality and did not yield TPU speedup at batch 1/core. The aggressive turbo cache therefore reproduced the throughput claim direction, but not acceptable quality on this 12-block stochastic SiT-B sampler.

Foresight policy: stale cache versus one-block LARA predictor

At matched policy decisions and 13–32% block reduction, the trained one-block DiT predictor improved FID by 6.99–57.06 points over stale activation reuse. The best predictor point was rho=0.98: FID 41.486 with 1.078x sampler-only speedup, versus the full reference FID 32.595. At the extreme 58.7% reduction, the predictor became unstable (FID 129.812).

This is an image-domain adaptation of the policy idea, not a direct reproduction of the text-to-video Foresight paper.

Layer-wise FastCache surrogate versus LARA skip-3

At approximately 42% matched block reduction:

Neither approximation preserved usable image quality at that aggressive route; LARA was less destructive but slower due to predictor and routing overhead.

Predictive cross-depth caching

On the official LARA 400k checkpoint, Selector-Conservative at 50 NFE reduced backbone block evaluations by 10%, reached FID 27.85 versus Full 27.23, and was effectively latency-neutral (38.62 versus 38.50 images/s). Static graph fusion remained necessary for compute reduction to become wall-clock acceleration.

Predictive recirculation / temporal activation feedback

The archive includes broad ADM-10K/50K sweeps for SLERP-based previous-step activation feedback across LARA-B, LARA-XL, official SiT, LayerSync, FLUX, and SD3 probes. One confirmed LARA-XL result used five cached mid-depth activations at alpha=0.02 and improved ADM-50K FID from 5.5593 to 5.3200 and sFID from 4.7890 to 4.6661, with a 1.57% median denoiser-latency increase.

Hardware and software

• —Google Cloud TPU VM, v5p-16, two workers / eight chips
• —Zone: us-central1-a
• —JAX 0.10.2 and libtpu 0.0.42.1 during the E-DiT experiments
• —TPU static DiT compilation required:

export LIBTPU_INIT_ARGS='--xla_tpu_rwb_fusion=false --xla_tpu_dot_dot_fusion=false'

• —OpenAI guided-diffusion ADM evaluator with 10,000 or 50,000 generated samples
• —SD VAE decode and ImageNet-256 reference statistics

Exact package listings and git state are in environment/.

Input provenance and licensing

This archive contains components with different upstream licenses; therefore the repository-level license is other. Users must follow each component's upstream license.

• —imagenet-vae-latents-ar-v2 and imagenet-vae-latents-train-v3 were downloaded from Kaggle and are marked CC0: Public Domain by their dataset publisher:
• —<https://www.kaggle.com/datasets/thaygiaodaysat/imagenet-vae-latents-ar-v2/versions/1>
• —<https://www.kaggle.com/datasets/thaygiaodaysat/imagenet-vae-latents-train-v3>
• —Inception-v3 PyTorch weights were downloaded from Kaggle; that variation is marked Apache-2.0: <https://www.kaggle.com/models/ctlcmleon/inception-v3/PyTorch/default/1>
• —SD VAE weights came from the Kaggle damtrunghieu/sdvae-ema Flax variation: <https://www.kaggle.com/models/damtrunghieu/sdvae-ema/Flax/default/1>. Consult the upstream model page and original Stable Diffusion VAE terms before redistribution or commercial use.
• —The ADM evaluator snapshot follows the upstream OpenAI guided-diffusion repository license: <https://github.com/openai/guided-diffusion>.
• —E-DiT references: paper and official repository.

The trained checkpoints are research artifacts derived from the listed inputs. No claim is made that a single permissive license covers every file.

Integrity and restoration

MANIFEST.tsv records relative path, byte size, and SHA-256 for every regular file in the archive. MANIFEST.sha256 is compatible with sha256sum -c from the repository root. To restore a subset:

hf download LamTNguyen/archive-TPUv5p-16-2026-08-26 \
  --repo-type dataset \
  --include 'artifacts/checkpoints/edit-sit-b-lara395k-static-hardste-cache-v1/**' \
  --local-dir ./archive

Large-folder upload state is resumable, and the final Hub file inventory was checked against the local manifest before VM cleanup.