reproducibility/build.py

"""Plain data-free NVFP4; deterministic source-supported FFN folding."""
import collections
import copy
import gc
import hashlib
import importlib.metadata
import json
import os
import re
import shutil
import time
from pathlib import Path

import torch
from huggingface_hub import HfApi, snapshot_download
from safetensors import safe_open
from safetensors.torch import save_file
from transformers import AutoModelForImageTextToText
import modelopt.torch.quantization as mtq
from modelopt.torch.quantization.nn import TensorQuantizer
from modelopt.torch.export import export_hf_checkpoint

SOURCE = 'Agnes-AI/Agnes-3.0-Flash'
REV = '24f712ce59379b54c4a141d2708c35daf5ff613b'
TARGET = 'ProCreations/Agnes-3.0-Flash-NVFP4'
ROOT = Path('/workspace/agnes')
OUTPUT = ROOT/'export'
API = HfApi()
SELECT = re.compile(r'^model\.language_model\.layers\.\d+\.(?:mlp\.(?:gate_proj|up_proj|down_proj)|global_attn\.(?:q_proj|k_proj|v_proj|o_proj))\.weight$')
REPORT = dict(source=SOURCE, source_revision=REV, target=TARGET,
    status='building', code_revision=os.environ['AGNES_CODE_REVISION'],
    modelopt_revision='5cae3940402f1ced98069a666b0bec72ec8b33b5',
    quantization=dict(method='plain NVFP4 max', training=False, calibration_examples=0,
        calibration_tokens=0, calibration_forward_calls=0, activation_global_scale=1.0,
        block_size=16, activation_block_scales='dynamic E4M3', kv_cache_quantized=False),
    folding=dict(main_width=17408, parallel_width=2048, folded_width=19456,
        gate_up_dimension=0, down_dimension=1, layers=72,
        reason='Match the upstream SGLang BF16 loader before quantization'),
    evaluation=dict(status='pending'))

def record():
    (ROOT/'quality_report.json').write_text(json.dumps(REPORT, indent=2)+'\n')
    API.upload_file(repo_id=TARGET, path_or_fileobj=str(ROOT/'quality_report.json'),
        path_in_repo='quality_report.json', commit_message='Record plain NVFP4 build checks')

def selected(name):
    return bool(SELECT.fullmatch(name))

def main():
    assert API.model_info(TARGET).private
    assert not any(n.endswith('.safetensors') for n in API.list_repo_files(TARGET))
    ROOT.mkdir(parents=True, exist_ok=True)
    torch.set_num_threads(8)
    torch.manual_seed(20260912)
    REPORT['versions'] = {p:importlib.metadata.version(p) for p in ['torch','transformers','nvidia-modelopt','accelerate']}
    REPORT['hardware'] = [torch.cuda.get_device_name(0)]
    record()
    from transformers.dynamic_module_utils import get_class_from_dynamic_module
    source = Path(snapshot_download(SOURCE, revision=REV, local_dir=ROOT/'source', allow_patterns=['*.py','config.json'], max_workers=8))
    cls = get_class_from_dynamic_module('modeling_agnes.AgnesForConditionalGeneration',source)
    assert cls.__name__ == 'AgnesForConditionalGeneration'
    print('AGNES_REMOTE_IMPLEMENTATION_IMPORT_PASS',flush=True)
    source = Path(snapshot_download(SOURCE, revision=REV, local_dir=ROOT/'source', max_workers=8))
    src_index = json.loads((source/'model.safetensors.index.json').read_text())
    src_map = src_index['weight_map']
    model = AutoModelForImageTextToText.from_pretrained(source, trust_remote_code=True,
        dtype=torch.bfloat16, device_map={'':0}, attn_implementation='sdpa').eval()
    assert all(p.device.type == 'cuda' for p in model.parameters())
    assert len(model.model.language_model.layers) == 72
    with torch.no_grad():
        for layer in model.model.language_model.layers:
            mlp = layer.mlp
            assert mlp.parallel_ffn is not None
            for name in ['gate_proj','up_proj','down_proj']:
                original, branch = getattr(mlp,name), getattr(mlp.parallel_ffn,name)
                dim = 1 if name == 'down_proj' else 0
                merged = torch.cat([original.weight,branch.weight], dim=dim)
                # Check both slices before discarding the separate branches.
                a,b = merged.split([original.weight.shape[dim],branch.weight.shape[dim]],dim=dim)
                assert torch.equal(a,original.weight) and torch.equal(b,branch.weight)
                replacement = torch.nn.Linear(merged.shape[1],merged.shape[0],bias=False,
                    device=merged.device,dtype=merged.dtype)
                replacement.weight = torch.nn.Parameter(merged,requires_grad=False)
                setattr(mlp,name,replacement)
            mlp.parallel_ffn = None
    text = model.config.text_config
    text.agnes_original_intermediate_size = 17408
    text.agnes_original_parallel_ffn_intermediate_size = 2048
    text.intermediate_size = 19456
    text.parallel_ffn_intermediate_size = 0
    gc.collect(); torch.cuda.empty_cache()
    preset = copy.deepcopy(mtq.NVFP4_DEFAULT_CFG)
    weight = next(x['cfg'] for x in preset['quant_cfg'] if x.get('quantizer_name')=='*weight_quantizer')
    activation = next(x['cfg'] for x in preset['quant_cfg'] if x.get('quantizer_name')=='*input_quantizer')
    activation['constant_amax'] = 2688.0
    rules = [dict(quantizer_name='*',enable=False)]
    targets = [name for name,module in model.named_modules()
               if isinstance(module,torch.nn.Linear) and selected(name+'.weight')]
    assert len(targets) == 288, len(targets)
    for name in targets:
        rules += [dict(quantizer_name=name+'.weight_quantizer',cfg=weight),
                  dict(quantizer_name=name+'.input_quantizer',cfg=activation)]
    cfg = dict(quant_cfg=rules,algorithm=dict(method='max',layerwise=dict(enable=False),
        skip_forward_without_activation_calib=True))
    REPORT['quantization']['configuration'] = cfg
    def forbid(*args):
        REPORT['quantization']['calibration_forward_calls'] += 1
        raise RuntimeError('Data-free quantization may not execute calibration forwards')
    hook = model.register_forward_pre_hook(forbid)
    started = time.monotonic()
    with torch.inference_mode():
        # AgnesDeltaAttention ends in "Attention" and its Python module exposes
        # ALL_ATTENTION_FUNCTIONS, so ModelOpt's heuristic incorrectly wraps it
        # as a standard attention class. KV quantization is explicitly out of
        # scope: omit that optional registration callback for this conversion.
        from modelopt.torch.quantization.plugins.custom import CUSTOM_MODEL_PLUGINS
        from modelopt.torch.quantization.plugins.huggingface import register_hf_attentions_on_the_fly
        assert register_hf_attentions_on_the_fly in CUSTOM_MODEL_PLUGINS
        CUSTOM_MODEL_PLUGINS.remove(register_hf_attentions_on_the_fly)
        try:
            mtq.quantize(model,cfg,forward_loop=None)
        finally:
            CUSTOM_MODEL_PLUGINS.add(register_hf_attentions_on_the_fly)
        REPORT['quantization']['kv_attention_wrapper_registration']=False
        enabled = [(name,m) for name,m in model.named_modules()
                   if isinstance(m,TensorQuantizer) and m.is_enabled]
        assert len(enabled) == 576, len(enabled)
        assert all(name.rsplit('.',1)[0] in targets for name,m in enabled)
        assert all(getattr(m,'_amax',None) is not None for name,m in enabled if name.endswith('weight_quantizer'))
        REPORT['quantization']['weight_statistics_seconds'] = time.monotonic()-started
        hook.remove()
        REPORT['export_graph_inspection'] = dict(dummy_forward_calls=0,input_tokens=0,quantizers_disabled=True)
        def inspect_probe(module,args):
            assert not any(isinstance(m,TensorQuantizer) and (m.is_enabled or m._if_calib) for m in module.modules())
            assert args[0].shape == (1,2) and torch.equal(args[0],torch.ones_like(args[0]))
            REPORT['export_graph_inspection']['dummy_forward_calls'] += 1
            REPORT['export_graph_inspection']['input_tokens'] += args[0].numel()
        hook = model.register_forward_pre_hook(inspect_probe)
        export_hf_checkpoint(model,dtype=torch.bfloat16,export_dir=OUTPUT,max_shard_size='5GB')
        hook.remove()
    del model,enabled
    gc.collect(); torch.cuda.empty_cache()
    # The upstream Transformers class ignores MTP: retain every original MTP
    # tensor separately, without advertising untested speculative decoding.
    dst_index = json.loads((OUTPUT/'model.safetensors.index.json').read_text())
    mtp = {}
    for name,filename in src_map.items():
        if name.startswith('mtp.'):
            with safe_open(source/filename,framework='pt',device='cpu') as f:
                mtp[name] = f.get_tensor(name).clone()
    assert mtp and not any(name in dst_index['weight_map'] for name in mtp)
    save_file(mtp,str(OUTPUT/'model-mtp.safetensors'),metadata={'format':'pt'})
    dst_index['weight_map'].update({name:'model-mtp.safetensors' for name in mtp})
    dst_index['metadata']['total_size'] += sum(t.numel()*t.element_size() for t in mtp.values())
    (OUTPUT/'model.safetensors.index.json').write_text(json.dumps(dst_index,indent=2)+'\n')
    for file in source.rglob('*'):
        rel=file.relative_to(source)
        if file.is_file() and '.cache' not in rel.parts and not file.name.endswith('.safetensors') and file.name not in ['config.json','model.safetensors.index.json','README.md','.gitattributes']:
            (OUTPUT/rel).parent.mkdir(parents=True,exist_ok=True)
            shutil.copy2(file,OUTPUT/rel)
    # The upstream loader otherwise skips renaming when the branch width is 0.
    # This checkpoint has already folded its branches before quantization.
    for file in OUTPUT.glob('sglang_patch/*/sglang/srt/models/qwen3_5.py'):
        contents=file.read_text()
        before='    if width <= 0:\n        yield from weights\n        return\n'
        after='    if width <= 0:\n        for name, weight in weights:\n            yield name.replace(".delta_attn.", ".linear_attn.").replace(".global_attn.", ".self_attn."), weight\n        return\n'
        assert contents.count(before)==1
        file.write_text(contents.replace(before,after))
    # Both Transformers names and translated serving names must be excluded.
    for filename in ['config.json','hf_quant_config.json']:
        file=OUTPUT/filename
        obj=json.loads(file.read_text())
        quant=obj['quantization_config'] if filename=='config.json' else obj['quantization']
        field = 'ignore' if filename == 'config.json' else 'exclude_modules'
        excludes=list(quant.get(field,[]))
        excludes += ['lm_head','model.visual*','mtp*']
        for i in range(72):
            excludes += [f'model.language_model.layers.{i}.delta_attn*',
                         f'model.language_model.layers.{i}.linear_attn*']
        quant[field]=sorted(set(excludes))
        file.write_text(json.dumps(obj,indent=2)+'\n')
    preserved=0; preserved_count=0
    for filename in sorted(set(src_map.values())):
        with safe_open(source/filename,framework='pt',device='cpu') as sf:
            for name in sf.keys():
                if selected(name) or '.mlp.parallel_ffn.' in name: continue
                assert name in dst_index['weight_map'],name
                with safe_open(OUTPUT/dst_index['weight_map'][name],framework='pt',device='cpu') as df:
                    a,b=sf.get_tensor(name),df.get_tensor(name)
                    assert a.dtype==b.dtype and torch.equal(a,b),name
                    preserved += a.numel()*a.element_size(); preserved_count+=1
    counts=collections.Counter(); packed=0; inputs=0
    for filename in sorted(set(dst_index['weight_map'].values())):
        with safe_open(OUTPUT/filename,framework='pt',device='cpu') as f:
            for name in f.keys():
                tensor=f.get_tensor(name)
                counts[str(tensor.dtype)]+=tensor.numel()*tensor.element_size()
                if selected(name):
                    assert tensor.dtype==torch.uint8,name
                    packed+=1
                if 'scale' in name:
                    assert torch.isfinite(tensor.float()).all() and (tensor.float()>0).all(),name
                if name.endswith('input_scale'):
                    assert torch.equal(tensor,torch.ones_like(tensor)),name
                    inputs+=1
    assert packed==len(targets)==288 and inputs==288,(packed,inputs)
    assert REPORT['quantization']['calibration_forward_calls']==0
    REPORT['export']=dict(source_tensor_bytes=src_index['metadata']['total_size'],
        exported_weight_bytes=sum((OUTPUT/f).stat().st_size for f in set(dst_index['weight_map'].values())),
        tensor_bytes_by_dtype=dict(counts),packed_linear_weights=packed,
        bf16_preserved_bytes=preserved,bf16_preserved_tensors=preserved_count,
        all_unquantized_tensors_bitwise_equal=True,mtp_tensors_preserved=len(mtp),
        activation_global_scales_exactly_one=inputs)
    REPORT['status']='packed_export_verified_pending_native_evaluation'
    record()
    shutil.copy2(ROOT/'quality_report.json',OUTPUT/'quality_report.json')
    API.upload_folder(repo_id=TARGET,folder_path=OUTPUT,
        commit_message='Upload verified plain NVFP4 Agnes Preview checkpoint')
    print('AGNES_BUILD_COMPLETE '+json.dumps(REPORT['export']),flush=True)

if __name__=='__main__':
    try: main()
    except Exception:
        REPORT['status']='build_failed'
        record()
        raise