|
|
|
@ -33,22 +33,6 @@ loss_fn_for_chatglm_model = lambda x: torch.nn.functional.mse_loss(
|
|
|
|
|
)
|
|
|
|
|
loss_fn = lambda x: x["loss"]
|
|
|
|
|
|
|
|
|
|
config = AutoConfig.from_pretrained(
|
|
|
|
|
"THUDM/chatglm2-6b",
|
|
|
|
|
trust_remote_code=True,
|
|
|
|
|
num_layers=2,
|
|
|
|
|
padded_vocab_size=65024,
|
|
|
|
|
hidden_size=64,
|
|
|
|
|
ffn_hidden_size=214,
|
|
|
|
|
num_attention_heads=8,
|
|
|
|
|
kv_channels=16,
|
|
|
|
|
rmsnorm=True,
|
|
|
|
|
original_rope=True,
|
|
|
|
|
use_cache=True,
|
|
|
|
|
multi_query_attention=False,
|
|
|
|
|
torch_dtype=torch.float32,
|
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
infer_config = AutoConfig.from_pretrained(
|
|
|
|
|
"THUDM/chatglm2-6b",
|
|
|
|
@ -68,6 +52,21 @@ infer_config = AutoConfig.from_pretrained(
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def init_chatglm():
|
|
|
|
|
config = AutoConfig.from_pretrained(
|
|
|
|
|
"THUDM/chatglm2-6b",
|
|
|
|
|
trust_remote_code=True,
|
|
|
|
|
num_layers=2,
|
|
|
|
|
padded_vocab_size=65024,
|
|
|
|
|
hidden_size=64,
|
|
|
|
|
ffn_hidden_size=214,
|
|
|
|
|
num_attention_heads=8,
|
|
|
|
|
kv_channels=16,
|
|
|
|
|
rmsnorm=True,
|
|
|
|
|
original_rope=True,
|
|
|
|
|
use_cache=True,
|
|
|
|
|
multi_query_attention=False,
|
|
|
|
|
torch_dtype=torch.float32,
|
|
|
|
|
)
|
|
|
|
|
model = AutoModelForCausalLM.from_config(config, empty_init=False, trust_remote_code=True)
|
|
|
|
|
for m in model.modules():
|
|
|
|
|
if m.__class__.__name__ == "RMSNorm":
|
|
|
|
|