Llama::load 和 forward 会 panic

cfg = AutoConfig.from_pretrained(model_dir, num_hidden_layers=31)
m = AutoModelForCausalLM.from_pretrained(model_dir, config=cfg, dtype=torch.float32)
print("31 layers loaded:", len(m.model.layers))
try:
  AutoModelForCausalLM.from_pretrained(model_dir, attn_implementation="flash_attention_2")
except Exception as e:
  print(type(e).__name__, str(e)[:200])

[transformers] LlamaForCausalLM LOAD REPORT from: ...
Key                                             | Status  | 
------------------------------------------------+---------+-
model.layers.30.self_attn.o_proj.weight         | MISSING | 
model.layers.30.self_attn.q_proj.weight         | MISSING | 
...
model.layers.30.mlp.gate_proj.weight            | MISSING | 

Notes:
- MISSING:	those params were newly initialized because missing from the checkpoint. Consider training on your downstream task.
31 layers loaded: 31
ImportError FlashAttention2 has been toggled on, but it cannot be used due to the following error: the package for FlashAttention2 doesn't seem to be installed.

far = model(ids[:, :1], position_ids=torch.tensor([[8192]])).logits

position 8192: (1, 1, 49152) finite True

let mut raw: serde_json::Value = serde_json::from_slice(&std::fs::read(format!("{MODEL}/config.json"))?)?;
if mode.starts_with("layer31") {
  raw["num_hidden_layers"] = 31.into();
}
let cfg: LlamaConfig = serde_json::from_value(raw)?;
let cfg = cfg.into_config(mode == "flash");
let model = Llama::load(vb, &cfg)?;
println!("loaded");
let mut cache = Cache::new(true, DType::F32, &cfg, &dev)?;
let (ids, pos): (&[u32], usize) = if mode.starts_with("far") { (&PROMPT[..1], 8192) } else { (&PROMPT, 0) };
let l = model.forward(&Tensor::new(ids, &dev)?.unsqueeze(0)?, pos, &mut cache)?;

== layer31

thread 'main' (9129703) panicked at .../candle-transformers-0.11.0/src/models/llama.rs:524:75:
called `Result::unwrap()` on an `Err` value: cannot find tensor model.layers.30.self_attn.q_proj.weight
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
exit 101
== flash
loaded

thread 'main' (9129725) panicked at .../candle-transformers-0.11.0/src/models/llama.rs:258:5:
not implemented: compile with '--features flash-attn'
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
exit 101
== far
loaded
Error: narrow invalid args start + len > dim_len: [8192, 32], dim: 0, start: 8192, len:1
exit 1

        let blocks: Vec<_> = (0..cfg.num_hidden_layers)
            .map(|i| Block::load(vb.pp(format!("model.layers.{i}")), cfg).unwrap())
            .collect();

#[cfg(not(feature = "flash-attn"))]
fn flash_attn(_: &Tensor, _: &Tensor, _: &Tensor, _: f32, _: bool) -> Result<Tensor> {
    unimplemented!("compile with '--features flash-attn'")
}

fn check_layers(vb: &VarBuilder, cfg: &Config) -> Result<()> {
  for i in 0..cfg.num_hidden_layers {
    for w in ["self_attn.q_proj", "self_attn.k_proj", "self_attn.v_proj", "self_attn.o_proj",
              "mlp.gate_proj", "mlp.up_proj", "mlp.down_proj", "input_layernorm", "post_attention_layernorm"] {
      let name = format!("model.layers.{i}.{w}.weight");
      if !vb.contains_tensor(&name) {
        bail!("config has {} layers, file has no {name}", cfg.num_hidden_layers);
      }
    }
  }
  Ok(())
}

== layer31-check
Error: config has 31 layers, file has no model.layers.30.self_attn.q_proj.weight
exit 1

let r = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| Llama::load(vb, &cfg)));
r.map_err(|p| anyhow::anyhow!("Llama::load panicked: {:?}", p.downcast_ref::<String>()))??

thread 'main' (9187338) panicked at .../candle-transformers-0.11.0/src/models/llama.rs:524:75:
called `Result::unwrap()` on an `Err` value: cannot find tensor model.layers.30.self_attn.q_proj.weight
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
Error: Llama::load panicked: Some("called `Result::unwrap()` on an `Err` value: cannot find tensor model.layers.30.self_attn.q_proj.weight")
exit 1

if pos + ids.len() > cfg.max_position_embeddings {
  bail!("position {} past max_position_embeddings {}", pos + ids.len() - 1, cfg.max_position_embeddings);
}

== far-check
loaded
Error: position 8192 past max_position_embeddings 8192
exit 1