When tranining the RWKV, it report "backward error"

[lxianl455]

When tranining the RWKV, it report "backward error"

[RUFFY-369]

Hi @lxianl455 please put your model in train mode with model.train() before performing a backprop

model = RwkvForCausalLM.from_pretrained("RWKV/rwkv-4-169m-pile", use_cache=True)
tokenizer.pad_token = tokenizer.eos_token
model.train()
.
.
.

Cheers!

[RUFFY-369]

@lxianl455 okay so if you want to to use the RWKV model just for inference or in default eval() mode and don't want to put it in train mode then modify your code to this, the error will go away:

tokenizer = AutoTokenizer.from_pretrained("RWKV/rwkv-4-169m-pile")
model = RwkvForCausalLM.from_pretrained("RWKV/rwkv-4-169m-pile", use_cache=False)
.
.
.

give use_cache a False value because during caching, the model stores intermediate results to speed up computation and it can interfere with gradient computation and also use_cache is not used for training or gradient computation.

Cheers!

[amyeroberts]

cc @ArthurZucker

[lxianl455]

Hi @lxianl455 please put your model in train mode with model.train() before performing a backprop

model = RwkvForCausalLM.from_pretrained("RWKV/rwkv-4-169m-pile", use_cache=True)
tokenizer.pad_token = tokenizer.eos_token
model.train()
.
.
.

Cheers!

Actually, I want to combine RWKV block with some other modules to predict time series information. I am not using the whole RWKV model, but only its blocks. In this scenario, the Transformer Trainer cannot be used. How can I solve this backward error?

[lxianl455]

Actually, the function I want is to work like LSTM. When training, LSTM can take in the init state for initialization and can also return the ending state afterwards. https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html

I think the use_cache in the RWKV code is not like the use_cache in other models. Its actual function is like asking for the return of LSTM cell state and hidden state. Here is the code (line 300 in modeling_rwkv.py ):

rwkv, layer_state = rwkv_linear_attention( self.time_decay, self.time_first, key, value, state=layer_state, return_state=use_cache, )

[ArthurZucker]

The recurrentGemma model implements something more in the lines of RNN (so close to LSTM) if you are looking for an equivalent)