rixty
/
KoboldAI-Client
mirror of https://github.com/KoboldAI/KoboldAI-Client.git
1
1
Fork 0
Code Issues Projects Releases Wiki Activity

Merge pull request #128 from VE-FORBRYDERNE/neox 

TPU support for HF GPT-NeoX model
This commit is contained in:
henk717 2022-06-22 01:39:53 +02:00 committed by GitHub
parent 75bc472a9f 8c594c6869
commit 5f9a116052
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 38 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
aiserver.py
View File

@ -976,6 +976,8 @@ if(vars.model not in ["InferKit", "Colab", "OAI", "GooseAI" , "ReadOnly", "TPUMe
if(vars.model_type == "opt"):
vars.badwordsids = vars.badwordsids_opt
if(vars.model_type == "neox"):
vars.badwordsids = vars.badwordsids_neox
if(not vars.use_colab_tpu and vars.model not in ["InferKit", "Colab", "OAI", "GooseAI" , "ReadOnly", "TPUMeshTransformerGPTJ", "TPUMeshTransformerGPTNeoX"]):
loadmodelsettings()
@ -1882,8 +1884,10 @@ else:
if vars.model in ("TPUMeshTransformerGPTJ", "TPUMeshTransformerGPTNeoX") and (not vars.custmodpth or not os.path.isdir(vars.custmodpth)):
raise FileNotFoundError(f"The specified model path {repr(vars.custmodpth)} is not the path to a valid folder")
import tpu_mtj_backend
if(vars.model == "TPUMeshTransformerGPTNeoX" or vars.model_type == "opt"):
if(vars.model_type == "opt"):
tpu_mtj_backend.pad_token_id = 1
elif(vars.model == "TPUMeshTransformerGPTNeoX" or vars.model_type == "neox"):
tpu_mtj_backend.pad_token_id = 2
tpu_mtj_backend.vars = vars
tpu_mtj_backend.warper_callback = tpumtjgenerate_warper_callback
tpu_mtj_backend.stopping_callback = tpumtjgenerate_stopping_callback

33
maps/gpt_neox.json Normal file
View File

@ -0,0 +1,33 @@
{
"mtj_compat": "neox",
"mtj_pe": "neox_rotary",
"mtj_config_map": {
"pe_rotary_pct": ["rotary_pct", 0.25],
"d_model": "hidden_size",
"n_heads": "num_attention_heads",
"layers": "num_hidden_layers"
},
"static_weights": {
"gpt_neox.embed_in.weight": {"mtj": {"module": "embedding_shard/~/linear", "param": "w", "transforms": ["no_transpose", "vocab_pad"]}},
"gpt_neox.final_layer_norm.weight": {"mtj": {"module": "projection_shard/~/replicated_layer_norm", "param": "scale"}},
"gpt_neox.final_layer_norm.bias": {"mtj": {"module": "projection_shard/~/replicated_layer_norm", "param": "offset"}},
"embed_out.weight": {"mtj": {"module": "projection_shard/~/linear", "param": "w", "transforms": ["vocab_pad"]}}
},
"layer_weights": {
"gpt_neox.layers.{layer}.attention.bias": {},
"gpt_neox.layers.{layer}.attention.masked_bias": {},
"gpt_neox.layers.{layer}.attention.rotary_emb.inv_freq": {},
"gpt_neox.layers.{layer}.attention.query_key_value.weight": {"mtj": {"module": "layer_{layer}/~/combined_qkv", "param": "w"}},
"gpt_neox.layers.{layer}.attention.query_key_value.bias": {"mtj": {"module": "layer_{layer}/~/combined_qkv", "param": "b"}},
"gpt_neox.layers.{layer}.attention.dense.weight": {"mtj": {"module": "layer_{layer}/~/linear_3", "param": "w"}},
"gpt_neox.layers.{layer}.attention.dense.bias": {"mtj": {"module": "layer_{layer}/~/linear_3", "param": "b", "transforms": ["divide_by_shards"]}},
"gpt_neox.layers.{layer}.mlp.dense_h_to_4h.weight": {"mtj": {"module": "layer_{layer}/~/linear_4", "param": "w"}},
"gpt_neox.layers.{layer}.mlp.dense_h_to_4h.bias": {"mtj": {"module": "layer_{layer}/~/linear_4", "param": "b"}},
"gpt_neox.layers.{layer}.mlp.dense_4h_to_h.weight": {"mtj": {"module": "layer_{layer}/~/linear_5", "param": "w"}},
"gpt_neox.layers.{layer}.mlp.dense_4h_to_h.bias": {"mtj": {"module": "layer_{layer}/~/linear_5", "param": "b", "transforms": ["divide_by_shards"]}},
"gpt_neox.layers.{layer}.input_layernorm.weight": {"mtj": {"module": "layer_{layer}/~/replicated_layer_norm", "param": "scale"}},
"gpt_neox.layers.{layer}.input_layernorm.bias": {"mtj": {"module": "layer_{layer}/~/replicated_layer_norm", "param": "offset"}},
"gpt_neox.layers.{layer}.post_attention_layernorm.weight": {"mtj": {"module": "layer_{layer}/~/replicated_layer_norm_1", "param": "scale"}},
"gpt_neox.layers.{layer}.post_attention_layernorm.bias": {"mtj": {"module": "layer_{layer}/~/replicated_layer_norm_1", "param": "offset"}}
}
}