Views
No views yet
1import torch
2from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
3
4model_id = "tiny-random/granite-4.0-h"
5tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
6model = AutoModelForCausalLM.from_pretrained(
7 model_id,
8 torch_dtype=torch.bfloat16,
9 trust_remote_code=True,
10)
11pipe = pipeline('text-generation', model=model, tokenizer=tokenizer, trust_remote_code=True)
12print(pipe('Write an article about Artificial Intelligence.'))1import json
2from pathlib import Path
3
4import accelerate
5import torch
6from huggingface_hub import file_exists, hf_hub_download
7from transformers import (
8 AutoConfig,
9 AutoModelForCausalLM,
10 AutoTokenizer,
11 GenerationConfig,
12 set_seed,
13)
14
15source_model_id = "ibm-granite/granite-4.0-h-small"
16save_folder = "/tmp/tiny-random/granite-4.0-h"
17
18processor = AutoTokenizer.from_pretrained(source_model_id)
19processor.save_pretrained(save_folder)
20
21with open(hf_hub_download(source_model_id, filename='config.json', repo_type='model'), 'r', encoding='utf-8') as f:
22 config_json = json.load(f)
23
24config_json['hidden_size'] = 32
25config_json['intermediate_size'] = 128
26config_json['layer_types'] = ['mamba', 'attention']
27config_json.update({
28 'mamba_expand': int(4096 / 32 * 2),
29})
30config_json['num_attention_heads'] = 2
31config_json['shared_intermediate_size'] = 128
32config_json['num_hidden_layers'] = 2
33config_json['num_key_value_heads'] = 2
34config_json['tie_word_embeddings'] = True
35
36with open(f"{save_folder}/config.json", "w", encoding='utf-8') as f:
37 json.dump(config_json, f, indent=2)
38
39config = AutoConfig.from_pretrained(
40 save_folder,
41 trust_remote_code=True,
42)
43print(config)
44
45torch.set_default_dtype(torch.bfloat16)
46model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)
47torch.set_default_dtype(torch.float32)
48
49if file_exists(filename="generation_config.json", repo_id=source_model_id, repo_type='model'):
50 model.generation_config = GenerationConfig.from_pretrained(
51 source_model_id, trust_remote_code=True,
52 )
53set_seed(42)
54model = model.cpu()
55with torch.no_grad():
56 for name, p in sorted(model.named_parameters()):
57 torch.nn.init.normal_(p, 0, 0.1)
58 print(name, p.shape)
59model.save_pretrained(save_folder)
60print(model)1GraniteMoeHybridForCausalLM(
2 (model): GraniteMoeHybridModel(
3 (embed_tokens): Embedding(100352, 32, padding_idx=100256)
4 (layers): ModuleList(
5 (0): GraniteMoeHybridDecoderLayer(
6 (block_sparse_moe): GraniteMoeHybridMoE(
7 (activation): SiLU()
8 (input_linear): GraniteMoeHybridParallelExperts()
9 (output_linear): GraniteMoeHybridParallelExperts()
10 (router): GraniteMoeHybridTopKGating(
11 (layer): Linear(in_features=32, out_features=72, bias=False)
12 )
13 )
14 (input_layernorm): GraniteMoeHybridRMSNorm((32,), eps=1e-05)
15 (post_attention_layernorm): GraniteMoeHybridRMSNorm((32,), eps=1e-05)
16 (shared_mlp): GraniteMoeHybridMLP(
17 (activation): SiLU()
18 (input_linear): Linear(in_features=32, out_features=256, bias=False)
19 (output_linear): Linear(in_features=128, out_features=32, bias=False)
20 )
21 (mamba): GraniteMoeHybridMambaLayer(
22 (act): SiLU()
23 (conv1d): Conv1d(8448, 8448, kernel_size=(4,), stride=(1,), padding=(3,), groups=8448)
24 (in_proj): Linear(in_features=32, out_features=16768, bias=False)
25 (norm): GraniteMoeHybridRMSNormGated()
26 (out_proj): Linear(in_features=8192, out_features=32, bias=False)
27 )
28 )
29 (1): GraniteMoeHybridDecoderLayer(
30 (block_sparse_moe): GraniteMoeHybridMoE(
31 (activation): SiLU()
32 (input_linear): GraniteMoeHybridParallelExperts()
33 (output_linear): GraniteMoeHybridParallelExperts()
34 (router): GraniteMoeHybridTopKGating(
35 (layer): Linear(in_features=32, out_features=72, bias=False)
36 )
37 )
38 (input_layernorm): GraniteMoeHybridRMSNorm((32,), eps=1e-05)
39 (post_attention_layernorm): GraniteMoeHybridRMSNorm((32,), eps=1e-05)
40 (shared_mlp): GraniteMoeHybridMLP(
41 (activation): SiLU()
42 (input_linear): Linear(in_features=32, out_features=256, bias=False)
43 (output_linear): Linear(in_features=128, out_features=32, bias=False)
44 )
45 (self_attn): GraniteMoeHybridAttention(
46 (q_proj): Linear(in_features=32, out_features=32, bias=False)
47 (k_proj): Linear(in_features=32, out_features=32, bias=False)
48 (v_proj): Linear(in_features=32, out_features=32, bias=False)
49 (o_proj): Linear(in_features=32, out_features=32, bias=False)
50 )
51 )
52 )
53 (norm): GraniteMoeHybridRMSNorm((32,), eps=1e-05)
54 )
55 (lm_head): Linear(in_features=32, out_features=100352, bias=False)
56)