Add bf16 mixed-precision inference option (--bf16)

[Chgara (chgara)]

Summary

Adds an inference option that runs the ViT encoder in bf16 while keeping the neck, heads and (for MoGe-3) the sparse refiner in fp32 — the same fine-grained mixed precision policy train_moge3.py --precision mixed_bf16 uses — with the encoder weights stored in bf16.

• MoGeModel.enable_mixed_precision(dtype, cast_encoder_weights=False) (v2/v3): new flag that also casts the encoder weights to dtype.
• MoGeModel.dtype (v2/v3) now reads the neck instead of the first parameter, so infer() keeps the input image and the uv grid in fp32 under mixed precision (as in training) and the refiner keeps receiving fp32 features.
• moge infer --bf16 (v2/v3 only, mutually exclusive with --fp16).

Motivation

The current --fp16 path for v3 keeps fp32 weights and wraps the whole forward in fp16 autocast. Two consequences:

1. Peak memory goes up, not down. Autocast caches the fp16 copies of fp32 weights for the duration of the forward pass, so moge-3-vitg peaks at roughly 5.0 GB (fp32 weights) + 2.5 GB (fp16 copies) of weight memory alone. The fp32 default is 5.0 GB. This PR's option is ≈ 2.7 GB (1.14B encoder params × 2 B + 0.11B × 4 B), and there is nothing to cache because the encoder weights already are bf16.
2. The heads and the refiner run in fp16 under --fp16, which they never saw in training; the encoder was trained under bf16 autocast, so bf16 for the encoder and fp32 for everything else is the policy the checkpoints were trained with.

Weight memory (GB = 10⁹ bytes):

	moge-3-vitg	moge-3-vitl	moge-2-vitl-normal
fp32 (default)	5.00	1.48	1.32
--fp16 (fp32 weights + autocast cache, peak)	~7.5	~2.2	n/a (.half())
--bf16 (this PR)	2.73	0.87	0.78

Numerically the matmuls see the same bf16 weights as under plain autocast; only the few fp32-policy ops (layer norms, positional-embedding interpolation) now read bf16-rounded parameters.

Why the dtype property change

infer() does image.to(self.dtype). With the encoder stored in bf16, next(self.parameters()) is a bf16 encoder parameter, so the image and the uv grid would silently become bf16 (training used fp32 for both), and for v3 torch.concat([features, uv]) would stay bf16 and mismatch the fp32 refiner. Reading the neck's dtype gives fp32 under mixed precision, fp16 after .half(), and fp32 for a plain model, i.e. it is unchanged for all existing code paths (self.dtype is only used in infer()).

Validation

• CPU check on a small randomly initialised MoGe-2 model (autocast hooks redirected to cpu for the test): encoder parameters bf16, all other modules fp32, model.dtype == torch.float32, infer() runs end to end with fp32 outputs of the same shapes; median relative depth deviation vs. the fp32 run 2.5e-4 (bf16-level noise).
• moge infer --help renders the new option; --bf16 with --version v1 or together with --fp16 raises a UsageError.
• The autocast policy itself is the one already exercised by train_moge3.py --precision mixed_bf16; the only new runtime code is self.encoder.to(dtype) and the neck-based dtype property.

[Chgara (chgara)]

@microsoft-github-policy-service agree company="Speridlabs"