Skip to content

Add native BF16 TRELLIS.2 inference - #1636

Open
for-the-juan wants to merge 3 commits into
ModelTC:mainfrom
for-the-juan:feat/trellis2-bf16
Open

for-the-juan wants to merge 3 commits into
ModelTC:mainfrom
for-the-juan:feat/trellis2-bf16

Conversation

@for-the-juan

Copy link
Copy Markdown

Summary

Split from #1628 at reviewer request. Add native BF16 TRELLIS.2-4B inference through LightX2V's runner registry (i23d). This PR targets main independently and requires no other model PR.

  • Native 512 dense/sparse flow networks, Euler/CFG sampling, input encoders, structure/shape/texture decoders and geometry GLB export.
  • Sorted integer neighbor lookup for portable sparse convolution, with an exact CPU hash reference. Dense and sparse linear adapters are loaded independently.
  • Registered BF16 deformable projection for BiRefNet: FP32 coordinate interpolation followed by BF16 sampled-feature/weight GEMMs. Replacement affects model instances, without globally patching torchvision.
  • Model-specific preparation tool, pinned public revisions, byte-checked input manifest, dependency pins and reproduction guide at scripts/bench/trellis2/README.md.

Model graphs, algorithms and operator selection are separate. Platform setup remains in existing lightx2v_platform; new model code has no vendor imports or branches. Shared dense adapters and measurement code are identical commits in the two split PRs, so they remain one implementation after merging. The other PR is #1635. Central runner/config registration entries may need rebasing after either PR merges; local composition was checked with both models present.

Validation

  • 21 model/shared numerical tests passed without exposing any accelerator; all 27 tests passed after composing the two branches locally.
  • pre-commit run --all-files passed.
  • Independent wheel build passed, with this model's license included and the other reconstruction model absent.
  • Post-split single-card MLU regression: two fixed requests / four A/B runs, zero failures, exact agreement of every returned tensor, and BF16 operand audits passed. Networking was disabled; the only source mount was this independent checkout. Dependency images were reused from the previously validated vendor-runtime builds.

These short runs validate the split; they are not a new complete performance measurement. The original complete frozen-cohort measurements are documented with their source revision (c20b99d) in the guide. The shared event/wall boundaries, alternating A/B order, seeds, input hashes and reductions are preserved; rerun the documented full suite for this checkout's timings.

Reproduction and scope

Run from the repository root:

python -m scripts.bench.trellis2.prepare_assets weights --output weights/trellis2
python -m scripts.bench.trellis2.benchmark \
  --model-path weights/trellis2 --input-root benchmark_inputs \
  --output benchmark_outputs/trellis2 --audit-bf16

Prepare the frozen inputs first using the guide. Expose one idle device per process. Matrix weights, large matrix/conv operands and attention Q/K/V are BF16; normalization, phase/coordinate arithmetic, accumulation and geometry may be FP32. Coordinates remain integers.

The formal cohort is 40 common-success GSO objects, seed 42 + order, 512, three × 12 steps, excluding the first pair for 39 observations. Events record 36 sample_once calls per request, reduced within each request then across requests. Wall time includes geometry export and texture voxels. This cohort does not estimate full-dataset success rate.

Only 512 and geometry GLB are included; no cascade/1024/1536 or baked textures. Some segment_reduce operations fall back to CPU in the tested torch_mlu version. The earlier experimental sparse kernel indexing, missed BF16 decoder storage and FP32 deformable path are corrected. Historical mesh metrics do not apply; full-dataset ground-truth quality has not been reevaluated for the corrected graph.

Weights, input data and the vendor runtime/driver remain external prerequisites; no private experiment checkout or patch is needed. Pinned BiRefNet Python files are public model assets, downloaded alongside its weights and loaded locally with trust_remote_code=True. The complete public GSO archive download was not repeated in this validation; archive/member hashes are enforced by the preparation command. CUDA/A100 and MetaX performance, offload, quantization and distributed inference have not been validated. Full-dataset GT metrics have not been rerun after splitting. Upstream CI status should be checked separately from local validation.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant