没有 no_grad

model = AutoModelForCausalLM.from_pretrained(model_dir, dtype=torch.float32, attn_implementation="eager").eval()
print("params requires_grad:", all(p.requires_grad for p in model.parameters()))
ids = torch.tensor([[504, 3575, 282, 4649, 314] * 102 + [504, 3575]])  # 512 个 token
if sys.argv[1] == "nograd":
  with torch.no_grad():
    logits = model(ids).logits
else:
  logits = model(ids).logits
print(f"logits.requires_grad {logits.requires_grad} grad_fn {type(logits.grad_fn).__name__}")

params requires_grad: True
logits.requires_grad True grad_fn UnsafeViewBackward0
          2071068672  maximum resident set size
params requires_grad: True
logits.requires_grad False grad_fn NoneType
          1143259136  maximum resident set size

let vb = unsafe { VarBuilder::from_mmaped_safetensors(&[&path], DType::F32, &dev)? };
let w = vb.get((49152, 576), "model.embed_tokens.weight")?;
println!("embed_tokens.weight is_variable {} track_op {}", w.is_variable(), w.track_op());
let model = Llama::load(vb, &cfg)?;
let ids = Tensor::new(&[504u32, 3575, 282, 4649, 314], &dev)?.unsqueeze(0)?;
let mut cache = Cache::new(false, DType::F32, &cfg, &dev)?;
let logits = model.forward(&ids, 0, &mut cache)?;
println!("logits track_op {} sorted_nodes {}", logits.track_op(), logits.sorted_nodes().len());

embed_tokens.weight is_variable false track_op false
logits track_op false sorted_nodes 0

x = torch.rand(2048, 2048) * 2 - 1
x.requires_grad_(sys.argv[1] == "grad")
a = x
for _ in range(20):
  a = a.tanh()

let x = Tensor::rand(-1f32, 1., (2048, 2048), &dev)?;
let x = if mode == "chainvar" { Var::from_tensor(&x)?.as_tensor().clone() } else { x };
let mut a = x;
for _ in 0..20 {
  a = a.tanh()?;
}

chainvar-drop5 round 4: track_op true sum -68.49649
           359464960  maximum resident set size

let varmap = VarMap::new();
let l = linear(3, 2, VarBuilder::from_varmap(&varmap, DType::F32, &dev).pp("fc"))?;
println!("train model output track_op {}", l.forward(&x)?.track_op());
let detached: HashMap<String, Tensor> =
  varmap.data().lock()?.iter().map(|(k, v)| (k.clone(), v.as_tensor().detach())).collect();
let eval = linear(3, 2, VarBuilder::from_tensors(detached, DType::F32, &dev).pp("fc"))?;
let before = eval.forward(&x)?;
println!("eval twin output track_op {}", before.track_op());
let mut opt = SGD::new(varmap.all_vars(), 0.1)?;
opt.backward_step(&l.forward(&x)?.sum_all()?)?;
println!("after step: train {:?} eval twin {:?}", l.forward(&x)?.to_vec2::<f32>()?, eval.forward(&x)?.to_vec2::<f32>()?);

train model output track_op true
eval twin output track_op false
before step [[-4.9460845, 2.523826]]
after step: train [[-6.446085, 1.0238259]] eval twin [[-6.446085, 1.0238259]]

p = torch.tensor([1.0, 2.0], requires_grad=True)
alias, det, cl, dc = p, p.detach(), p.clone(), p.detach().clone()
opt = torch.optim.SGD([p], lr=0.5)
(p ** 2).sum().backward()
opt.step()

p after step       [0.0, 0.0]
p.detach()         [0.0, 0.0]
p.clone()          [1.0, 2.0] requires_grad True
p.detach().clone() [1.0, 2.0]

let p = Var::new(&[1f32, 2.], &dev)?;
let snap_clone = p.as_tensor().clone();
let snap_detach = p.as_tensor().detach();
let snap_copy = p.as_tensor().copy()?;
let snap_contig = (p.as_tensor() * 1.0)?;
let mut opt = SGD::new(vec![p.clone()], 0.5)?;
opt.backward_step(&p.sqr()?.sum_all()?)?;

x = torch.tensor([[1.0, 2.0, 3.0, 4.0]], requires_grad=True)
g = torch.ones(4, requires_grad=True)
c = torch.tensor([[1.0, -1.0, 2.0, 0.5]])
(F.rms_norm(x, (4,), g, 1e-5) * c).sum().backward()

let cases: [(&str, Tensor); 4] = [
  ("ops::rms_norm", ops::rms_norm(&x, &g, 1e-5)?),
  ("ops::rms_norm_slow", ops::rms_norm_slow(&x, &g, 1e-5)?),
  ("ops::softmax_last_dim", ops::softmax_last_dim(&x)?),
  ("ops::softmax(.., D::Minus1)", ops::softmax(&x, D::Minus1)?),
];
for (name, y) in cases {
  let grads = (y * &c)?.sum_all()?.backward()?;
  // ... 打印 x 和 g 的梯度
}

F.rms_norm   x.grad [0.2799470126628876, -0.5355503559112549, 0.4746928811073303, -0.15823042392730713] g.grad [0.36514812707901, -0.73029625415802, 2.1908888816833496, 0.73029625415802]
softmax(-1)  x.grad [0.008314814418554306, -0.1516866236925125, 0.2983214259147644, -0.1549495905637741] g.grad None

ops::rms_norm                x.grad None g.grad is_some false
ops::rms_norm_slow           x.grad Some([0.27994698, -0.53555036, 0.47469288, -0.15823045]) g.grad is_some true
ops::softmax_last_dim        x.grad None g.grad is_some false
ops::softmax(.., D::Minus1)  x.grad Some([0.0083148135, -0.15168662, 0.29832146, -0.15494965]) g.grad is_some false

for (name, t) in candle_core::safetensors::load(&path, &dev)? {
  let v = Var::from_tensor(&t.to_dtype(DType::F32)?)?;
  map.insert(name, v.as_tensor().clone());
  vars.push(v);
}
let model = Llama::load(VarBuilder::from_tensors(map, DType::F32, &dev), &cfg)?;
let logits = model.forward(&ids, 0, &mut cache)?;
let nodes = logits.sorted_nodes();
let g = logits.sum_all()?.backward()?;

logits track_op true sorted_nodes 3 (vars among them 1)
272 vars, 1 with a gradient

p after step   [0.0, 0.0]
clone()        [0.0, 0.0]
detach()       [0.0, 0.0]
copy()         [1.0, 2.0]
(p * 1.0)      [1.0, 2.0]

p after step   [0.0, 0.0]
clone()        [0.0, 0.0]
detach()       [0.0, 0.0]
copy()         [0.0, 0.0]
(p * 1.0)      [1.0, 2.0]