Training checkpoint saved in fp32 (2x size vs pretrained), causes inference OOM

[xpwindells]

After fine-tuning SenseNova-Vision-7B-MoT on a custom panoptic segmentation dataset with FSDP HYBRID_SHARD (8 GPUs), the saved checkpoint weights (ema.safetensors / model.safetensors) are 54.4 GB (fp32), while the original pretrained checkpoint is only 27.2 GB (bf16). This 2x size increase causes OOM during single-GPU inference on an 80GB card.

Environment
Model: SenseNova-Vision-7B-MoT
Training: 8× A100 80GB, FSDP HYBRID_SHARD (num_shard=4, num_replicate=4)
Checkpoint: results/.../checkpoint/0002000/ema.safetensors
Inference: single GPU (A100 80GB), SenseNovaVisionModel from inference/sensenova_vision.py

Line 105-112: EMA model save

with FSDP.state_dict_type(
ema_model,
StateDictType.FULL_STATE_DICT,
FullStateDictConfig(rank0_only=True, offload_to_cpu=True),
):
ema_state_dict = ema_model.state_dict()
if dist.get_rank() == 0:
save_file(ema_state_dict, os.path.join(save_path, "ema.safetensors")) # ← fp32!

Line 114-121: Model save (same issue)

with FSDP.state_dict_type(
model,
StateDictType.FULL_STATE_DICT,
FullStateDictConfig(rank0_only=True, offload_to_cpu=True),
):
model_state_dict = model.state_dict()
if dist.get_rank() == 0:
save_file(model_state_dict, os.path.join(save_path, "model.safetensors")) # ← fp32!

[Robin970822]

Thanks for the detailed report — the checkpoint-size comparison and environment information are very helpful.

The FP32 training checkpoint is expected. FSDP uses BF16 for training computation, but we keep the saved model and EMA states in full precision so that resuming training preserves the expected precision. For inference, we recommend creating a separate BF16 checkpoint instead of replacing the FP32 training checkpoint.

Automatic BF16 conversion during inference

SenseNovaVisionModel defaults to dtype="bf16" and resolves it to torch.bfloat16:

SenseNova-Vision/inference/sensenova_vision.py

Lines 245 to 256 in 884d515

	def __init__(
	self,
	model_path: str = "sensenova/SenseNova-Vision-7B-MoT",
	checkpoint_name: str = "ema.safetensors",
	dtype: str = "bf16",
	device: str = "auto",
	max_mem_per_gpu: str = "80GiB",
	offload_folder: str = "offload",
	seed: Optional[int] = 42,
	download_revision: str = "main",
	download_local_files_only: bool = False,
	):

The resolved dtype is then passed to load_checkpoint_and_dispatch:

SenseNova-Vision/inference/sensenova_vision.py

Lines 498 to 520 in 884d515

	def _load_model_weights(self, model):
	device_map = self._infer_device_map(model)
	if self.dtype in [
	"bf16",
	"bfloat16",
	"fp16",
	"float16",
	"half",
	"fp32",
	"float32",
	"auto",
	"none",
	]:
	model = load_checkpoint_and_dispatch(
	model,
	checkpoint=self.checkpoint_path,
	device_map=device_map,
	offload_folder=self.offload_folder,
	offload_buffers=True,
	dtype=self.torch_dtype,
	force_hooks=True,
	).eval()

. Inference computation also runs under BF16 autocast:

SenseNova-Vision/inference/inferencer.py

Lines 273 to 307 in 884d515

	@torch.no_grad()
	def interleave_inference(
	self,
	input_lists: List[Union[str, Image.Image]],
	think=False,
	understanding_output=False,
	caption=False,
	output_multiple_vae=False,
	output_raw_tensor=False,
	return_preprocessed_input=False,
	max_think_token_n=1000,
	do_sample=False,
	text_temperature=0.3,
	cfg_text_scale=3.0,
	cfg_img_scale=1.5,
	cfg_interval=[0.4, 1.0],
	timestep_shift=3.0,
	num_timesteps=50,
	cfg_renorm_min=0.0,
	cfg_renorm_type="global",
	image_shapes=(1024, 1024),
	vae_transform=None,
	vit_transform=None,
	) -> List[Union[str, Image.Image]]:
	output_list = []
	gen_context = self.init_gen_context()
	cfg_text_context = deepcopy(gen_context)
	cfg_img_context = deepcopy(gen_context)
	preprocessed_inputs = [] if return_preprocessed_input else None
	vae_transform = vae_transform or self.vae_transform
	vit_transform = vit_transform or self.vit_transform
	with torch.autocast(device_type="cuda", enabled=True, dtype=torch.bfloat16):

Therefore, the final in-memory model should use BF16. However, when the source checkpoint is FP32, the automatic conversion still happens during loading and may temporarily increase peak GPU memory.

Explicitly converting the checkpoint to BF16

To avoid the load-time FP32-to-BF16 conversion peak, please create a separate BF16 inference checkpoint first.

You can download the conversion helper here:

convert_safetensors_bf16.py

View conversion script

import argparse
import os

import torch
from safetensors import safe_open
from safetensors.torch import save_file


FLOAT_DTYPES = (
    torch.float32,
    torch.float16,
    torch.float64,
    torch.bfloat16,
)


def parse_args():
    parser = argparse.ArgumentParser(
        description="Convert floating tensors in a safetensors checkpoint to bf16."
    )
    parser.add_argument("--input", required=True, help="Input .safetensors path")
    parser.add_argument("--output", required=True, help="Output .safetensors path")
    return parser.parse_args()


def main():
    args = parse_args()

    tensors = {}
    total_in = 0
    total_out = 0

    with safe_open(args.input, framework="pt", device="cpu") as f:
        metadata = f.metadata()

        for key in f.keys():
            tensor = f.get_tensor(key)
            total_in += tensor.numel() * tensor.element_size()

            if tensor.dtype in FLOAT_DTYPES:
                tensor = tensor.to(torch.bfloat16)

            total_out += tensor.numel() * tensor.element_size()
            tensors[key] = tensor

    output_dir = os.path.dirname(args.output)
    if output_dir:
        os.makedirs(output_dir, exist_ok=True)

    save_file(tensors, args.output, metadata=metadata)

    print("Saved:", args.output)
    print(f"Approx bytes in : {total_in / 1024**3:.2f} GB")
    print(f"Approx bytes out: {total_out / 1024**3:.2f} GB")
    print(
        "Actual file size:",
        os.path.getsize(args.output) / 1024**3,
        "GB",
    )


if __name__ == "__main__":
    main()

Run it with:

python convert_safetensors_bf16.py \
  --input /path/to/checkpoint/ema.safetensors \
  --output /path/to/checkpoint/ema_bf16.safetensors

Then load the converted checkpoint:

model = SenseNovaVisionModel(
    model_path="/path/to/checkpoint",
    checkpoint_name="ema_bf16.safetensors",
    dtype="bf16",
)

The conversion runs on CPU and preserves the original FP32 checkpoint for training resume.

If the BF16 checkpoint still produces an OOM, could you share the complete traceback and confirm whether it happens:

• During model loading.
• During automatic dtype conversion/dispatch.
• During VAE encode.
• During text/image output generation or VAE decode.

Please also include the last log line before the OOM and the printed device_map.