Views
No views yet
pip install -U triton git+https://github.com/triton-lang/triton.git@main#subdirectory=python/triton_kernelsvllm serve tiny-random/gpt-oss-mxfp41import torch
2from transformers import pipeline
3
4model_id = "tiny-random/gpt-oss-mxfp4"
5
6pipe = pipeline(
7 "text-generation",
8 model=model_id,
9 torch_dtype='auto',
10 device_map="cuda",
11)
12
13messages = [
14 {"role": "user", "content": "Explain quantum mechanics clearly and concisely."},
15]
16
17outputs = pipe(
18 messages,
19 max_new_tokens=16,
20)
21print(outputs[0]["generated_text"][-1])1import json
2
3import safetensors
4import torch
5from huggingface_hub import hf_hub_download
6from transformers import (
7 AutoConfig,
8 AutoModelForCausalLM,
9 AutoProcessor,
10 AutoTokenizer,
11 GenerationConfig,
12 GptOssForCausalLM,
13 pipeline,
14 set_seed,
15)
16
17source_model_id = "openai/gpt-oss-120b"
18save_folder = "/tmp/tiny-random/gpt-oss-mxfp4"
19
20processor = AutoProcessor.from_pretrained(source_model_id)
21processor.save_pretrained(save_folder)
22
23with open(hf_hub_download(source_model_id, filename='config.json', repo_type='model'), 'r') as f:
24 config_json = json.load(f)
25config_json.update({
26 "head_dim": 32,
27 "hidden_size": 32, # required by Mxfp4GptOssExperts codes
28 "intermediate_size": 64,
29 "layer_types": ["sliding_attention", "full_attention"],
30 "num_attention_heads": 2,
31 "num_hidden_layers": 2,
32 "num_key_value_heads": 1,
33 "num_local_experts": 32,
34 "tie_word_embeddings": True,
35})
36quantization_config = config_json['quantization_config']
37del config_json['quantization_config']
38with open(f"{save_folder}/config.json", "w", encoding='utf-8') as f:
39 json.dump(config_json, f, indent=2)
40
41config = AutoConfig.from_pretrained(save_folder)
42print(config)
43torch.set_default_dtype(torch.bfloat16)
44model = AutoModelForCausalLM.from_config(config, torch_dtype=torch.bfloat16)
45torch.set_default_dtype(torch.float32)
46model.generation_config = GenerationConfig.from_pretrained(
47 source_model_id, trust_remote_code=True,
48)
49set_seed(42)
50with torch.no_grad():
51 for name, p in sorted(model.named_parameters()):
52 torch.nn.init.normal_(p, 0, 0.1)
53 print(name, p.shape)
54model.save_pretrained(save_folder)
55
56# mxfp4
57state_dict = model.cpu().state_dict()
58del state_dict['lm_head.weight']
59for i in range(len(model.model.layers)):
60 del state_dict[f'model.layers.{i}.mlp.experts.down_proj']
61 del state_dict[f'model.layers.{i}.mlp.experts.gate_up_proj']
62 state_dict[f'model.layers.{i}.mlp.experts.down_proj_blocks'] = torch.randint(0, 255, size=(
63 config.num_local_experts, config.hidden_size, config.intermediate_size // 32, 16), dtype=torch.uint8
64 )
65 state_dict[f'model.layers.{i}.mlp.experts.down_proj_scales'] = torch.randint(0, 4, size=(
66 config.num_local_experts, config.hidden_size, config.intermediate_size // 32), dtype=torch.uint8
67 )
68 state_dict[f'model.layers.{i}.mlp.experts.gate_up_proj_blocks'] = torch.randint(0, 255, size=(
69 config.num_local_experts, 2 * config.intermediate_size, config.hidden_size // 32, 16), dtype=torch.uint8
70 )
71 state_dict[f'model.layers.{i}.mlp.experts.gate_up_proj_scales'] = torch.randint(0, 4, size=(
72 config.num_local_experts, 2 * config.intermediate_size, config.hidden_size // 32), dtype=torch.uint8
73 )
74safetensors.torch.save_file(state_dict, f"{save_folder}/model.safetensors")
75
76# from unittest.mock import Mock
77# from transformers.quantizers.auto import AutoHfQuantizer
78# from transformers.quantizers.quantizer_mxfp4 import Mxfp4HfQuantizer
79# _get_device_capability = torch.cuda.get_device_capability
80# torch.cuda.get_device_capability = Mock(return_value=(9, 0))
81# set_seed(42)
82# bf16_state_dict = model.cpu().state_dict()
83# model = AutoModelForCausalLM.from_pretrained(save_folder, torch_dtype=torch.bfloat16, quantization_config=quantization_config)
84# for i in range(len(model.model.layers)):
85# model.model.layers[i].mlp.experts.down_proj_bottom_pad = 0
86# model.model.layers[i].mlp.experts.down_proj_right_pad = 0
87# hf_quantizer: Mxfp4HfQuantizer = AutoHfQuantizer.from_config(quantization_config)
88# hf_quantizer.pre_quantized = False
89# ffn_keys = ['model.layers.0.mlp.experts.down_proj', 'model.layers.0.mlp.experts.gate_up_proj',
90# 'model.layers.1.mlp.experts.down_proj', 'model.layers.1.mlp.experts.gate_up_proj']
91# for key in ffn_keys:
92# hf_quantizer.create_quantized_param(model, bf16_state_dict[key], key, "cuda", bf16_state_dict)
93# print('down_proj', model.model.layers[0].mlp.experts.down_proj)
94# print('down_proj_blocks', model.model.layers[0].mlp.experts.down_proj_blocks)
95# state_dict = model.state_dict()
96# del state_dict['lm_head.weight']
97# for key in ffn_keys:
98# del state_dict[key]
99# for k, v in state_dict.items():
100# if str(v.device) == 'meta':
101# print(k, v.device, v.shape)
102
103# safetensors.torch.save_file(state_dict, f"{save_folder}/model.safetensors")
104with open(f"{save_folder}/config.json", "r", encoding='utf-8') as f:
105 config = json.load(f)
106config['quantization_config'] = quantization_config
107with open(f"{save_folder}/config.json", "w", encoding='utf-8') as f:
108 json.dump(config, f, indent=2)
109# torch.cuda.get_device_capability = _get_device_capability