Please help to shoot the error while loading locally trained model to chat.py

[shwangdev]

Hi,

I was training the model locally from scratch.

 python train.py --load_model  --wandb  --proj_dir out --data_file ../data/enwik8 --data_type utf-8 --vocab_size 0 --ctx_len 512 --epoch_steps 5000 --epoch_count 500 --epoch_begin 0 --epoch_save 5 --micro_bsz 12 --n_layer 6 --n_embd 512 --pre_ffn 0 --head_qk 0 --lr_init 8e-4 --lr_final 1e-5 --warmup_steps 0 --beta1 0.9 --beta2 0.99 --adam_eps 1e-8 --accelerator gpu --devices 1 --precision tf32 --strategy ddp_find_unused_parameters_false --grad_cp 0

and changed settings in chat.py

args.FLOAT_MODE = "fp32" # fp32 (good for CPU) // fp16 (recommended for GPU) // bf16 (less accurate)
args.vocab_size = 50277
args.head_qk = 0
args.pre_ffn = 0
args.grad_cp = 0
args.my_pos_emb = 0

# args.MODEL_NAME = '/fsx/BlinkDL/HF-MODEL/rwkv-4-pile-14b/RWKV-4-Pile-14B-20230108-5170'
args.MODEL_NAME = './out/rwkv-40'

args.n_layer = 6 # 40
args.n_embd = 512 # 5120
args.ctx_len = 512 # 1024

Getting an error:

(py38) ➜  RWKV-v4neo git:(main) ✗ python chat.py
Loading...

RWKV_HEAD_QK_DIM 0 RWKV_JIT_ON 1

loading... ./out/rwkv-40
emb.weight                               float32    cpu
blocks.0.ln1.weight                      float32    cuda:0
blocks.0.ln1.bias                        float32    cuda:0
blocks.0.ln2.weight                      float32    cuda:0
blocks.0.ln2.bias                        float32    cuda:0
blocks.0.ln0.weight                      float32    cuda:0
blocks.0.ln0.bias                        float32    cuda:0
blocks.0.att.time_decay                  float32    cuda:0
blocks.0.att.time_first                  float32    cuda:0
blocks.0.att.time_mix_k                  float32    cuda:0
blocks.0.att.time_mix_v                  float32    cuda:0
blocks.0.att.time_mix_r                  float32    cuda:0
blocks.0.att.key.weight                  float32    cuda:0
blocks.0.att.value.weight                float32    cuda:0
blocks.0.att.receptance.weight           float32    cuda:0
blocks.0.att.output.weight               float32    cuda:0
blocks.0.ffn.time_mix_k                  float32    cuda:0
blocks.0.ffn.time_mix_r                  float32    cuda:0
blocks.0.ffn.key.weight                  float32    cuda:0
blocks.0.ffn.receptance.weight           float32    cuda:0
blocks.0.ffn.value.weight                float32    cuda:0
..........................................................................................
ln_out.weight                            float32    cuda:0
ln_out.bias                              float32    cuda:0
head.weight                              float32    cuda:0

Run prompt...
Traceback (most recent call last):
  File "chat.py", line 193, in <module>
    out = run_rnn(tokenizer.tokenizer.encode(init_prompt))
  File "chat.py", line 163, in run_rnn
    current_state = model.forward(model_tokens, current_state, preprocess_only = True)
  File "/mnt/d/workspace/RWKV-LM/RWKV-v4neo/src/model_run.py", line 200, in forward
    x = w.emb.weight[ctx[-1]]
IndexError: index 48656 is out of bounds for dimension 0 with size 6064

would please help to take look if I mistake something. Thanks.