Fine-Tuning NemotronOmni on CORD-v2 Receipts — End-to-End Guide

1
import json
2
from datasets import load_dataset
3
4
dataset = load_dataset("naver-clova-ix/cord-v2")
5
6
print(f"Train      : {len(dataset['train'])} samples")
7
print(f"Validation : {len(dataset['validation'])} samples")
8
print(f"Test       : {len(dataset['test'])} samples")
9
10
# Inspect a sample
11
ex = dataset["train"][0]
12
gt = json.loads(ex["ground_truth"])["gt_parse"]
13
print(f"\nGround-truth keys: {list(gt.keys())}")

1
import torch
2
import json
3
from transformers import AutoModel, AutoProcessor
4
from datasets import load_dataset
5
from nemo_automodel.components.datasets.vlm.utils import json2token
6
7
CKPT = "<checkpoint_dir>/LOWEST_VAL/model/consolidated"
8
9
# Load processor
10
processor = AutoProcessor.from_pretrained(CKPT, trust_remote_code=True)
11
tokenizer = processor.tokenizer
12
13
# `device_map` streams weights directly to GPU; skipping the AutoModel.from_config
14
# CPU-instantiation step saves ~5 min on the 30B v3 dump.
15
model = AutoModel.from_pretrained(
16
    CKPT, trust_remote_code=True, torch_dtype=torch.bfloat16,
17
    device_map={"": torch.cuda.current_device()},
18
)
19
20
# Reset RADIO's `summary_idxs` (non-persistent buffer; can be a meta tensor after load)
21
if hasattr(model, "vision_model") and hasattr(model.vision_model, "radio_model"):
22
    model.vision_model.radio_model.summary_idxs = None
23
24
model.eval()
25
26
# Load dataset
27
dataset = load_dataset("naver-clova-ix/cord-v2")
28
29
# v3 processor returns extra placeholder-expansion metadata that is NOT a generate() kwarg.
30
PROCESSOR_METADATA_KEYS = ("num_patches", "num_tokens", "imgs_sizes")
31
32
# Run inference on the first 5 validation samples
33
for i in range(5):
34
    sample = dataset["validation"][i]
35
    image = sample["image"].convert("RGB")
36
    gt = json.loads(sample["ground_truth"])["gt_parse"]
37
    gt_text = json2token(gt, sort_json_key=True)
38
39
    # Build prompt — enable_thinking=False for structured output
40
    messages = [{"role": "user", "content": "<image>\nDescribe this image."}]
41
    text = tokenizer.apply_chat_template(
42
        messages, tokenize=False,
43
        add_generation_prompt=True, enable_thinking=False,
44
    )
45
    inputs = processor(text=text, images=[image], return_tensors="pt")
46
    for k in PROCESSOR_METADATA_KEYS:
47
        inputs.pop(k, None)
48
    inputs = {k: v.to("cuda") if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
49
50
    with torch.no_grad():
51
        output_ids = model.generate(**inputs, max_new_tokens=1024, do_sample=False)
52
53
    generated = tokenizer.decode(
54
        output_ids[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True,
55
    ).strip()
56
57
    print(f"\n=== Sample {i} ===")
58
    print(f"Ground truth: {gt_text}")
59
    print(f"Prediction:   {generated}")

1
import json, re
2
import torch
3
from pathlib import Path
4
from safetensors import safe_open
5
from transformers import AutoModel, AutoProcessor
6
7
BASE    = "<path_to_nemotron_omni_v3>"
8
ADAPTER = "<ckpt_dir>/LOWEST_VAL/model"
9
10
# Load base directly to GPU. Skip AutoModel.from_config — instantiating a 30B
11
# model on CPU just to read the class type adds 5+ minutes.
12
processor = AutoProcessor.from_pretrained(BASE, trust_remote_code=True)
13
model = AutoModel.from_pretrained(
14
    BASE, trust_remote_code=True, dtype=torch.bfloat16,
15
    device_map={"": torch.cuda.current_device()},
16
)
17
if hasattr(model, "vision_model") and hasattr(model.vision_model, "radio_model"):
18
    model.vision_model.radio_model.summary_idxs = None
19
20
# Wrapper -> HF base FQN translation. vision_projector.* targets are listed in
21
# adapter_config.json but no tensors are saved for them, so we just skip those.
22
def translate(fqn):
23
    if fqn.startswith("language_model.model."):
24
        return "language_model.backbone." + fqn[len("language_model.model."):]
25
    return None
26
27
cfg   = json.loads((Path(ADAPTER) / "adapter_config.json").read_text())
28
scale = cfg["lora_alpha"] / cfg["r"]
29
30
pairs = {}
31
with safe_open(str(Path(ADAPTER) / "adapter_model.safetensors"), framework="pt") as f:
32
    for k in f.keys():
33
        m = re.match(r"^base_model\.model\.(.+)\.lora_(A|B)\.weight$", k)
34
        if m:
35
            pairs.setdefault(m.group(1), {})[m.group(2)] = f.get_tensor(k)
36
37
modules = dict(model.named_modules())
38
for wrapper_fqn, ab in pairs.items():
39
    hf_fqn = translate(wrapper_fqn)
40
    if hf_fqn is None or hf_fqn not in modules:
41
        continue
42
    W = modules[hf_fqn].weight
43
    A = ab["A"].to(device=W.device, dtype=torch.float32)
44
    B = ab["B"].to(device=W.device, dtype=torch.float32)
45
    with torch.no_grad():
46
        W.add_(((B @ A) * scale).to(W.dtype))
47
48
model.eval()
49
# ... then run the same generate() loop as in the SFT example above.