Upload folder using huggingface_hub

fdc6474 verified 3 months ago

5.49 kB

#!/usr/bin/env python3
"""Byte-exact VRAM budget for GLM-5.2 hybrid quant on 4x RTX PRO 6000 (96GB, SM120).
Model constants from zai-org/GLM-5.2 config.json (verified against the
lukealonso/GLM-5.2-NVFP4 safetensors index: predicted 450.99 GB vs actual
450.93 GB total_size).
"""
GB = 1e9
GiB = 2**30

# --- architecture ---
H = 6144                 # hidden_size
MOE_I = 2048             # moe_intermediate_size
DENSE_I = 12288          # intermediate_size (dense layers 0-2)
VOCAB = 154880
N_LAYERS = 78            # main layers
N_MTP = 1                # layer 78
N_DENSE = 3              # layers 0,1,2
N_EXPERT_LAYERS = 75 + 1 # layers 3..77 + MTP layer
N_EXPERTS = 256
KV_LORA = 512
QK_ROPE = 64
Q_LORA = 2048
N_HEADS = 64
QK_NOPE = 192
V_HEAD = 256
IDX_HEADS, IDX_HDIM = 32, 128
N_FULL_INDEXERS = 22     # 21 main (IndexShare freq=4) + 1 MTP
CTX = 1_048_576

# --- per-component parameter counts ---
expert_params_down = MOE_I * H            # 12.58M
expert_params_gateup = 2 * H * MOE_I      # 25.17M
expert_params = expert_params_down + expert_params_gateup   # 37.75M
layer_expert_params = N_EXPERTS * expert_params             # 9.664B

attn_params = (H*Q_LORA + Q_LORA*(N_HEADS*(QK_NOPE+QK_ROPE))
               + H*(KV_LORA+QK_ROPE) + KV_LORA*(N_HEADS*(QK_NOPE+V_HEAD))
               + (N_HEADS*V_HEAD)*H)      # 165.03M per layer
indexer_params = Q_LORA*IDX_HEADS*IDX_HDIM + H*IDX_HDIM + H*IDX_HEADS  # 9.37M
shared_expert_params = expert_params
dense_mlp_params = 3 * H * DENSE_I
gate_params = H * N_EXPERTS
embed_params = VOCAB * H

TOTAL_LAYER_ENTRIES = N_LAYERS + N_MTP    # 79

nonexpert = {
    'attention':      attn_params * TOTAL_LAYER_ENTRIES,
    'indexers':       indexer_params * N_FULL_INDEXERS,
    'dense_mlp':      dense_mlp_params * N_DENSE,
    'shared_experts': shared_expert_params * N_EXPERT_LAYERS,
    'gates_fp32':     gate_params * N_EXPERT_LAYERS,      # router kept fp32
    'embed+lm_head':  2 * embed_params,
    'mtp_extras':     2 * H * H + 4 * H,                  # eh_proj etc
}

# --- bytes/param for each quant recipe (routed experts) ---
def gguf_bpw(down_bpw, gateup_bpw):
    return (expert_params_down*down_bpw + expert_params_gateup*gateup_bpw) / expert_params

RECIPES = {
    # NVFP4: 4b weight + fp8 scale per 16 = 4.5 bpw (+ negligible per-tensor scales)
    'nvfp4':  4.5,
    # unsloth UD-Q3_K_XL donor: down=IQ4_XS(4.25), gate/up=IQ3_XXS(3.0625)
    'ud_q3':  gguf_bpw(4.25, 3.0625),
    # unsloth UD-Q2_K_XL donor: down=IQ3_XXS(3.0625), gate/up=IQ2_XS(2.3125)
    'ud_q2':  gguf_bpw(3.0625, 2.3125),
}

def layer_expert_bytes(recipe):
    return layer_expert_params * RECIPES[recipe] / 8

# --- KV cache at 1M ---
kv_latent_per_tok = (KV_LORA + QK_ROPE) * TOTAL_LAYER_ENTRIES        # values (MTP has KV too)
def kv_bytes(bytes_per_val, scale_overhead=0.02):
    return kv_latent_per_tok * bytes_per_val * CTX * (1 + scale_overhead)
idx_cache = N_FULL_INDEXERS * IDX_HDIM * 1 * CTX   # fp8 indexer k-cache

# --- GPU budget ---
GPU_MEM = 97887 * 2**20        # RTX PRO 6000 Blackwell, nvidia-smi reported
N_GPU = 4
RESERVE_PER_GPU = 6 * GiB      # cuda ctx + nccl + cudagraphs + activations + frag

def report(attn_dtype_bytes=2, n4=0, n3=0, n2=0, kv='fp8'):
    assert n4+n3+n2 == 75
    ne = dict(nonexpert)
    scale = {'attention': attn_dtype_bytes/2, 'shared_experts': attn_dtype_bytes/2,
             'dense_mlp': attn_dtype_bytes/2}
    ne_bytes = 0
    for k, p in ne.items():
        b = p * (4 if k=='gates_fp32' else 2)
        b *= scale.get(k, 1)
        ne_bytes += b
    expert_bytes = (n4*layer_expert_bytes('nvfp4') + n3*layer_expert_bytes('ud_q3')
                    + n2*layer_expert_bytes('ud_q2') + layer_expert_bytes('nvfp4'))  # +MTP nvfp4
    weights = ne_bytes + expert_bytes
    kvb = kv_bytes(1 if kv=='fp8' else 0.5) + idx_cache
    total_need = weights + kvb
    total_have = N_GPU * (GPU_MEM - RESERVE_PER_GPU)
    print(f'--- attn={16*attn_dtype_bytes//2}bit  experts: {n4}xNVFP4 {n3}xUD-Q3 {n2}xUD-Q2 (+MTP nvfp4), KV={kv} ---')
    print(f'  non-expert weights: {ne_bytes/GB:8.2f} GB')
    print(f'  expert weights:     {expert_bytes/GB:8.2f} GB   (avg {expert_bytes*8/ (N_EXPERT_LAYERS*layer_expert_params):.2f} bpw)')
    print(f'  KV latent @1M:      {kv_bytes(1 if kv=="fp8" else 0.5)/GB:8.2f} GB')
    print(f'  indexer cache @1M:  {idx_cache/GB:8.2f} GB')
    print(f'  TOTAL NEED:         {total_need/GB:8.2f} GB = {total_need/GiB:.1f} GiB')
    print(f'  TOTAL HAVE (4x{GPU_MEM/GiB:.1f}GiB - 4x{RESERVE_PER_GPU/GiB:.0f}GiB): {total_have/GB:.2f} GB = {total_have/GiB:.1f} GiB')
    margin = total_have - total_need
    print(f'  MARGIN:             {margin/GB:+8.2f} GB   {"FITS" if margin>0 else "DOES NOT FIT"}')
    # PP=4 per-stage worst case: stage holding embed+lm_head + its layers
    return margin

print('kv latent values/token:', kv_latent_per_tok, '->', kv_latent_per_tok, 'B/tok fp8,',
      f'{kv_bytes(1)/GB:.1f} GB @1M fp8')
print('nonexpert param count:', sum(nonexpert.values())/1e9, 'B params')
print('checkpoint check: nvfp4-all =',
      (sum(v*2 for v in nonexpert.values()) + 2*nonexpert['gates_fp32'] +
       76*layer_expert_bytes('nvfp4'))/GB, 'GB vs actual 450.93')
print()
print('=== v1: BF16 attention (as shipped in NVFP4 repo) ===')
report(2, 10, 30, 35)
report(2, 8, 26, 41)
report(2, 6, 22, 47)
print()
print('=== v2: FP8 attention/shared/dense (unsloth uses Q8_0 there too) ===')
report(1, 14, 40, 21)
report(1, 12, 36, 27)
report(1, 10, 30, 35)

	#!/usr/bin/env python3
	"""Byte-exact VRAM budget for GLM-5.2 hybrid quant on 4x RTX PRO 6000 (96GB, SM120).

	Model constants from zai-org/GLM-5.2 config.json (verified against the
	lukealonso/GLM-5.2-NVFP4 safetensors index: predicted 450.99 GB vs actual
	450.93 GB total_size).
	"""
	GB = 1e9
	GiB = 2**30

	# --- architecture ---
	H = 6144 # hidden_size
	MOE_I = 2048 # moe_intermediate_size
	DENSE_I = 12288 # intermediate_size (dense layers 0-2)
	VOCAB = 154880
	N_LAYERS = 78 # main layers
	N_MTP = 1 # layer 78
	N_DENSE = 3 # layers 0,1,2
	N_EXPERT_LAYERS = 75 + 1 # layers 3..77 + MTP layer
	N_EXPERTS = 256
	KV_LORA = 512
	QK_ROPE = 64
	Q_LORA = 2048
	N_HEADS = 64
	QK_NOPE = 192
	V_HEAD = 256
	IDX_HEADS, IDX_HDIM = 32, 128
	N_FULL_INDEXERS = 22 # 21 main (IndexShare freq=4) + 1 MTP
	CTX = 1_048_576

	# --- per-component parameter counts ---
	expert_params_down = MOE_I * H # 12.58M
	expert_params_gateup = 2 * H * MOE_I # 25.17M
	expert_params = expert_params_down + expert_params_gateup # 37.75M
	layer_expert_params = N_EXPERTS * expert_params # 9.664B

	attn_params = (HQ_LORA + Q_LORA(N_HEADS*(QK_NOPE+QK_ROPE))
	+ H(KV_LORA+QK_ROPE) + KV_LORA(N_HEADS*(QK_NOPE+V_HEAD))
	+ (N_HEADSV_HEAD)H) # 165.03M per layer
	indexer_params = Q_LORAIDX_HEADSIDX_HDIM + HIDX_HDIM + HIDX_HEADS # 9.37M
	shared_expert_params = expert_params
	dense_mlp_params = 3 * H * DENSE_I
	gate_params = H * N_EXPERTS
	embed_params = VOCAB * H

	TOTAL_LAYER_ENTRIES = N_LAYERS + N_MTP # 79

	nonexpert = {
	'attention': attn_params * TOTAL_LAYER_ENTRIES,
	'indexers': indexer_params * N_FULL_INDEXERS,
	'dense_mlp': dense_mlp_params * N_DENSE,
	'shared_experts': shared_expert_params * N_EXPERT_LAYERS,
	'gates_fp32': gate_params * N_EXPERT_LAYERS, # router kept fp32
	'embed+lm_head': 2 * embed_params,
	'mtp_extras': 2 * H * H + 4 * H, # eh_proj etc
	}

	# --- bytes/param for each quant recipe (routed experts) ---
	def gguf_bpw(down_bpw, gateup_bpw):
	return (expert_params_downdown_bpw + expert_params_gateupgateup_bpw) / expert_params

	RECIPES = {
	# NVFP4: 4b weight + fp8 scale per 16 = 4.5 bpw (+ negligible per-tensor scales)
	'nvfp4': 4.5,
	# unsloth UD-Q3_K_XL donor: down=IQ4_XS(4.25), gate/up=IQ3_XXS(3.0625)
	'ud_q3': gguf_bpw(4.25, 3.0625),
	# unsloth UD-Q2_K_XL donor: down=IQ3_XXS(3.0625), gate/up=IQ2_XS(2.3125)
	'ud_q2': gguf_bpw(3.0625, 2.3125),
	}

	def layer_expert_bytes(recipe):
	return layer_expert_params * RECIPES[recipe] / 8

	# --- KV cache at 1M ---
	kv_latent_per_tok = (KV_LORA + QK_ROPE) * TOTAL_LAYER_ENTRIES # values (MTP has KV too)
	def kv_bytes(bytes_per_val, scale_overhead=0.02):
	return kv_latent_per_tok * bytes_per_val * CTX * (1 + scale_overhead)
	idx_cache = N_FULL_INDEXERS * IDX_HDIM * 1 * CTX # fp8 indexer k-cache

	# --- GPU budget ---
	GPU_MEM = 97887 * 2**20 # RTX PRO 6000 Blackwell, nvidia-smi reported
	N_GPU = 4
	RESERVE_PER_GPU = 6 * GiB # cuda ctx + nccl + cudagraphs + activations + frag

	def report(attn_dtype_bytes=2, n4=0, n3=0, n2=0, kv='fp8'):
	assert n4+n3+n2 == 75
	ne = dict(nonexpert)
	scale = {'attention': attn_dtype_bytes/2, 'shared_experts': attn_dtype_bytes/2,
	'dense_mlp': attn_dtype_bytes/2}
	ne_bytes = 0
	for k, p in ne.items():
	b = p * (4 if k=='gates_fp32' else 2)
	b *= scale.get(k, 1)
	ne_bytes += b
	expert_bytes = (n4layer_expert_bytes('nvfp4') + n3layer_expert_bytes('ud_q3')
	+ n2*layer_expert_bytes('ud_q2') + layer_expert_bytes('nvfp4')) # +MTP nvfp4
	weights = ne_bytes + expert_bytes
	kvb = kv_bytes(1 if kv=='fp8' else 0.5) + idx_cache
	total_need = weights + kvb
	total_have = N_GPU * (GPU_MEM - RESERVE_PER_GPU)
	print(f'--- attn={16*attn_dtype_bytes//2}bit experts: {n4}xNVFP4 {n3}xUD-Q3 {n2}xUD-Q2 (+MTP nvfp4), KV={kv} ---')
	print(f' non-expert weights: {ne_bytes/GB:8.2f} GB')
	print(f' expert weights: {expert_bytes/GB:8.2f} GB (avg {expert_bytes8/ (N_EXPERT_LAYERSlayer_expert_params):.2f} bpw)')
	print(f' KV latent @1M: {kv_bytes(1 if kv=="fp8" else 0.5)/GB:8.2f} GB')
	print(f' indexer cache @1M: {idx_cache/GB:8.2f} GB')
	print(f' TOTAL NEED: {total_need/GB:8.2f} GB = {total_need/GiB:.1f} GiB')
	print(f' TOTAL HAVE (4x{GPU_MEM/GiB:.1f}GiB - 4x{RESERVE_PER_GPU/GiB:.0f}GiB): {total_have/GB:.2f} GB = {total_have/GiB:.1f} GiB')
	margin = total_have - total_need
	print(f' MARGIN: {margin/GB:+8.2f} GB {"FITS" if margin>0 else "DOES NOT FIT"}')
	# PP=4 per-stage worst case: stage holding embed+lm_head + its layers
	return margin

	print('kv latent values/token:', kv_latent_per_tok, '->', kv_latent_per_tok, 'B/tok fp8,',
	f'{kv_bytes(1)/GB:.1f} GB @1M fp8')
	print('nonexpert param count:', sum(nonexpert.values())/1e9, 'B params')
	print('checkpoint check: nvfp4-all =',
	(sum(v2 for v in nonexpert.values()) + 2nonexpert['gates_fp32'] +
	76*layer_expert_bytes('nvfp4'))/GB, 'GB vs actual 450.93')
	print()
	print('=== v1: BF16 attention (as shipped in NVFP4 repo) ===')
	report(2, 10, 30, 35)
	report(2, 8, 26, 41)
	report(2, 6, 22, 47)
	print()
	print('=== v2: FP8 attention/shared/dense (unsloth uses Q8_0 there too) ===')
	report(1, 14, 40, 21)
	report(1, 12, 36, 27)
	report(1, 10, 30, 35)