Error with stories15M and stories110M

[deafTim]

When I use stories15M and stories110M I got an error.

  File "D:\_LLM_project\Development\gpt-fast\generate.py", line 114, in speculative_decode
    torch.cat([cur_token.view(1), draft_tokens]).view(1, -1),
RuntimeError: Tensors must have same number of dimensions: got 1 and 2

cur_token: tensor(6324, dtype=torch.int32)
cur_token shape: torch.Size([])
cur_token.view(1): tensor([6324], dtype=torch.int32)
cur_token.view(1,1): tensor([[6324]], dtype=torch.int32)
draft_tokens: tensor([[23717],
        [29889],
        [ 6324],
        [23717],
        [  338]], dtype=torch.int32)
draft_tokens shape: torch.Size([5, 1])

Please, help

[malfet]

I don't think anyone ever tried speculative decoding with tiny stories...

[deafTim]

I don't think anyone ever tried speculative decoding with tiny stories...

@malfet
I also think so, but could it be a problem? Isn't right that all should be the same for all the models from the model.py?

[deafTim]

I use this

torch.cat([cur_token.view(1,1), draft_tokens]).view(1, -1)

Instead of this

torch.cat([cur_token.view(1), draft_tokens]).view(1, -1)

But I got another error

p = draft_probs[torch.arange(0, speculate_k, device=device), draft_tokens]
IndexError: index 23717 is out of bounds for dimension 1 with size 1

[deafTim]

The issue is partially resolved—I adjusted some shapes to get it working with stories15M and stories110M models.
However, it doesn’t yet work for batch sizes greater than one (bs>1).
I’ll continue making code changes to support larger batch sizes.
In the meantime, is there an option to find an appropriate model around the 1B parameter size?

[malfet]

@deafTim https://huggingface.co/meta-llama/Llama-3.2-1B :)

[deafTim]

@malfet Thanks! I also need a smaller model for speculative decoding. The Stories model doesn’t seem to work very well.

[deafTim]

@malfet When I use meta-llama/Llama-3.2-1B
Can it be fixed?

RuntimeError: Error(s) in loading state_dict for Transformer:
	Missing key(s) in state_dict: "tok_embeddings.weight", "layers.0.attention.wqkv.weight", "layers.0.attention.wo.weight", "layers.0.feed_forward.w1.weight", "layers.0.feed_forward.w3.weight", "layers.0.feed_forward.w2.weight", "layers.0.ffn_norm.weight", "layers.0.attention_norm.weight", "layers.1.attention.wqkv.weight", "layers.1.attention.wo.weight", "layers.1.feed_forward.w1.weight", "layers.1.feed_forward.w3.weight", "layers.1.feed_forward.w2.weight", "layers.1.ffn_norm.weight", "layers.1.attention_norm.weight", "layers.2.attention.wqkv.weight", "layers.2.attention.wo.weight", "layers.2.feed_forward.w1.weight", "layers.2.feed_forward.w3.weight", "layers.2.feed_forward.w2.weight", "layers.2.ffn_norm.weight", "layers.2.attention_norm.weight", "layers.3.attention.wqkv.weight", "layers.3.attention.wo.weight", "layers.3.feed_forward.w1.weight", "layers.3.feed_forward.w3.weight", "layers.3.feed_forward.w2.weight", "layers.3.ffn_norm.weight", "layers.3.attention_norm.weight", "layers.4.attention.wqkv.weight", "layers.4.attention.wo.weight", "layers.4.feed_forward.w1.weight", "layers.4.feed_forward.w3.weight", "layers.4.feed_forward.w2.weight", "layers.4.ffn_norm.weight", "layers.4.attention_norm.weight", "layers.5.attention.wqkv.weight", "layers.5.attention.wo.weight", "layers.5.feed_forward.w1.weight", "layers.5.feed_forward.w3.weight", "layers.5.feed_forward.w2.weight", "layers.5.ffn_norm.weight", "layers.5.attention_norm.weight", "layers.6.attention.wqkv.weight", "layers.6.attention.wo.weight", "layers.6.feed_forward.w1.weight", "layers.6.feed_forward.w3.weight", "layers.6.feed_forward.w2.weight", "layers.6.ffn_norm.weight", "layers.6.attention_norm.weight", "layers.7.attention.wqkv.weight", "layers.7.attention.wo.weight", "layers.7.feed_forward.w1.weight", "layers.7.feed_forward.w3.weight", "layers.7.feed_forward.w2.weight", "layers.7.ffn_norm.weight", "layers.7.attention_norm.weight", "layers.8.attention.wqkv.weight", "layers.8.attention.wo.weight", "layers.8.feed_forward.w1.weight", "layers.8.feed_forward.w3.weight", "layers.8.feed_forward.w2.weight", "layers.8.ffn_norm.weight", "layers.8.attention_norm.weight", "layers.9.attention.wqkv.weight", "layers.9.attention.wo.weight", "layers.9.feed_forward.w1.weight", "layers.9.feed_forward.w3.weight", "layers.9.feed_forward.w2.weight", "layers.9.ffn_norm.weight", "layers.9.attention_norm.weight", "layers.10.attention.wqkv.weight", "layers.10.attention.wo.weight", "layers.10.feed_forward.w1.weight", "layers.10.feed_forward.w3.weight", "layers.10.feed_forward.w2.weight", "layers.10.ffn_norm.weight", "layers.10.attention_norm.weight", "layers.11.attention.wqkv.weight", "layers.11.attention.wo.weight", "layers.11.feed_forward.w1.weight", "layers.11.feed_forward.w3.weight", "layers.11.feed_forward.w2.weight", "layers.11.ffn_norm.weight", "layers.11.attention_norm.weight", "layers.12.attention.wqkv.weight", "layers.12.attention.wo.weight", "layers.12.feed_forward.w1.weight", "layers.12.feed_forward.w3.weight", "layers.12.feed_forward.w2.weight", "layers.12.ffn_norm.weight", "layers.12.attention_norm.weight", "layers.13.attention.wqkv.weight", "layers.13.attention.wo.weight", "layers.13.feed_forward.w1.weight", "layers.13.feed_forward.w3.weight", "layers.13.feed_forward.w2.weight", "layers.13.ffn_norm.weight", "layers.13.attention_norm.weight", "layers.14.attention.wqkv.weight", "layers.14.attention.wo.weight", "layers.14.feed_forward.w1.weight", "layers.14.feed_forward.w3.weight", "layers.14.feed_forward.w2.weight", "layers.14.ffn_norm.weight", "layers.14.attention_norm.weight", "layers.15.attention.wqkv.weight", "layers.15.attention.wo.weight", "layers.15.feed_forward.w1.weight", "layers.15.feed_forward.w3.weight", "layers.15.feed_forward.w2.weight", "layers.15.ffn_norm.weight", "layers.15.attention_norm.weight", "norm.weight", "output.weight". 
	Unexpected key(s) in state_dict: "model.embed_tokens.weight", "model.layers.0.input_layernorm.weight", "model.layers.0.mlp.down_proj.weight", "model.layers.0.mlp.gate_proj.weight", "model.layers.0.mlp.up_proj.weight", "model.layers.0.post_attention_layernorm.weight", "model.layers.0.self_attn.k_proj.weight", "model.layers.0.self_attn.o_proj.weight", "model.layers.0.self_attn.q_proj.weight", "model.layers.0.self_attn.v_proj.weight", "model.layers.1.input_layernorm.weight", "model.layers.1.mlp.down_proj.weight", "model.layers.1.mlp.gate_proj.weight", "model.layers.1.mlp.up_proj.weight", "model.layers.1.post_attention_layernorm.weight", "model.layers.1.self_attn.k_proj.weight", "model.layers.1.self_attn.o_proj.weight", "model.layers.1.self_attn.q_proj.weight", "model.layers.1.self_attn.v_proj.weight", "model.layers.10.input_layernorm.weight", "model.layers.10.mlp.down_proj.weight", "model.layers.10.mlp.gate_proj.weight", "model.layers.10.mlp.up_proj.weight", "model.layers.10.post_attention_layernorm.weight", "model.layers.10.self_attn.k_proj.weight", "model.layers.10.self_attn.o_proj.weight", "model.layers.10.self_attn.q_proj.weight", "model.layers.10.self_attn.v_proj.weight", "model.layers.11.input_layernorm.weight", "model.layers.11.mlp.down_proj.weight", "model.layers.11.mlp.gate_proj.weight", "model.layers.11.mlp.up_proj.weight", "model.layers.11.post_attention_layernorm.weight", "model.layers.11.self_attn.k_proj.weight", "model.layers.11.self_attn.o_proj.weight", "model.layers.11.self_attn.q_proj.weight", "model.layers.11.self_attn.v_proj.weight", "model.layers.12.input_layernorm.weight", "model.layers.12.mlp.down_proj.weight", "model.layers.12.mlp.gate_proj.weight", "model.layers.12.mlp.up_proj.weight", "model.layers.12.post_attention_layernorm.weight", "model.layers.12.self_attn.k_proj.weight", "model.layers.12.self_attn.o_proj.weight", "model.layers.12.self_attn.q_proj.weight", "model.layers.12.self_attn.v_proj.weight", "model.layers.13.input_layernorm.weight", "model.layers.13.mlp.down_proj.weight", "model.layers.13.mlp.gate_proj.weight", "model.layers.13.mlp.up_proj.weight", "model.layers.13.post_attention_layernorm.weight", "model.layers.13.self_attn.k_proj.weight", "model.layers.13.self_attn.o_proj.weight", "model.layers.13.self_attn.q_proj.weight", "model.layers.13.self_attn.v_proj.weight", "model.layers.14.input_layernorm.weight", "model.layers.14.mlp.down_proj.weight", "model.layers.14.mlp.gate_proj.weight", "model.layers.14.mlp.up_proj.weight", "model.layers.14.post_attention_layernorm.weight", "model.layers.14.self_attn.k_proj.weight", "model.layers.14.self_attn.o_proj.weight", "model.layers.14.self_attn.q_proj.weight", "model.layers.14.self_attn.v_proj.weight", "model.layers.15.input_layernorm.weight", "model.layers.15.mlp.down_proj.weight", "model.layers.15.mlp.gate_proj.weight", "model.layers.15.mlp.up_proj.weight", "model.layers.15.post_attention_layernorm.weight", "model.layers.15.self_attn.k_proj.weight", "model.layers.15.self_attn.o_proj.weight", "model.layers.15.self_attn.q_proj.weight", "model.layers.15.self_attn.v_proj.weight", "model.layers.2.input_layernorm.weight", "model.layers.2.mlp.down_proj.weight", "model.layers.2.mlp.gate_proj.weight", "model.layers.2.mlp.up_proj.weight", "model.layers.2.post_attention_layernorm.weight", "model.layers.2.self_attn.k_proj.weight", "model.layers.2.self_attn.o_proj.weight", "model.layers.2.self_attn.q_proj.weight", "model.layers.2.self_attn.v_proj.weight", "model.layers.3.input_layernorm.weight", "model.layers.3.mlp.down_proj.weight", "model.layers.3.mlp.gate_proj.weight", "model.layers.3.mlp.up_proj.weight", "model.layers.3.post_attention_layernorm.weight", "model.layers.3.self_attn.k_proj.weight", "model.layers.3.self_attn.o_proj.weight", "model.layers.3.self_attn.q_proj.weight", "model.layers.3.self_attn.v_proj.weight", "model.layers.4.input_layernorm.weight", "model.layers.4.mlp.down_proj.weight", "model.layers.4.mlp.gate_proj.weight", "model.layers.4.mlp.up_proj.weight", "model.layers.4.post_attention_layernorm.weight", "model.layers.4.self_attn.k_proj.weight", "model.layers.4.self_attn.o_proj.weight", "model.layers.4.self_attn.q_proj.weight", "model.layers.4.self_attn.v_proj.weight", "model.layers.5.input_layernorm.weight", "model.layers.5.mlp.down_proj.weight", "model.layers.5.mlp.gate_proj.weight", "model.layers.5.mlp.up_proj.weight", "model.layers.5.post_attention_layernorm.weight", "model.layers.5.self_attn.k_proj.weight", "model.layers.5.self_attn.o_proj.weight", "model.layers.5.self_attn.q_proj.weight", "model.layers.5.self_attn.v_proj.weight", "model.layers.6.input_layernorm.weight", "model.layers.6.mlp.down_proj.weight", "model.layers.6.mlp.gate_proj.weight", "model.layers.6.mlp.up_proj.weight", "model.layers.6.post_attention_layernorm.weight", "model.layers.6.self_attn.k_proj.weight", "model.layers.6.self_attn.o_proj.weight", "model.layers.6.self_attn.q_proj.weight", "model.layers.6.self_attn.v_proj.weight", "model.layers.7.input_layernorm.weight", "model.layers.7.mlp.down_proj.weight", "model.layers.7.mlp.gate_proj.weight", "model.layers.7.mlp.up_proj.weight", "model.layers.7.post_attention_layernorm.weight", "model.layers.7.self_attn.k_proj.weight", "model.layers.7.self_attn.o_proj.weight", "model.layers.7.self_attn.q_proj.weight", "model.layers.7.self_attn.v_proj.weight", "model.layers.8.input_layernorm.weight", "model.layers.8.mlp.down_proj.weight", "model.layers.8.mlp.gate_proj.weight", "model.layers.8.mlp.up_proj.weight", "model.layers.8.post_attention_layernorm.weight", "model.layers.8.self_attn.k_proj.weight", "model.layers.8.self_attn.o_proj.weight", "model.layers.8.self_attn.q_proj.weight", "model.layers.8.self_attn.v_proj.weight", "model.layers.9.input_layernorm.weight", "model.layers.9.mlp.down_proj.weight", "model.layers.9.mlp.gate_proj.weight", "model.layers.9.mlp.up_proj.weight", "model.layers.9.post_attention_layernorm.weight", "model.layers.9.self_attn.k_proj.weight", "model.layers.9.self_attn.o_proj.weight", "model.layers.9.self_attn.q_proj.weight", "model.layers.9.self_attn.v_proj.weight", "model.norm.weight". 

[deafTim]

It seems like this model has a different architecture. Is there a way to fix that?
@malfet