MEMIT and PMET torch.linalg.solve stuck

[Wonder1905]

Im trying to run knowledge editing using MEMIT and PMET , but the run always get stuck, after debugging I found out that the reason is : adj_k = torch.linalg.solve( hparams.mom2_update_weight * cov + layer_ks @ layer_ks.T, layer_ks, )

layer_ks and cov was calculated in double precision, I turned it into float, still didnt help. Am I the only one that is affected by this issue?

[XeeKee]

I can't seem to reproduce your issue. Could you provide your hyperparameters? Also, may I know what machine you are using? I'm able to run it successfully on an A800

[zxlzr]

Hi, have you solved your issue yet?

[Wonder1905]

Those are my hyperparameters: alg_name: "MEMIT" model_name: "meta-llama/Meta-Llama-3-8B" stats_dir: "./data/stats" device: 0 layers: [ 5, 6, 7, 8] clamp_norm_factor: 4 layer_selection: "all" fact_token: "subject_last" v_num_grad_steps: 20 v_lr: 5e-1 v_loss_layer: 20 v_weight_decay: 1e-3 kl_factor: 0.0625 mom2_adjustment: true mom2_update_weight: 1500 rewrite_module_tmp: "model.layers.{}.mlp.down_proj" layer_module_tmp: "model.layers.{}" mlp_module_tmp: "model.layers.{}.mlp" attn_module_tmp: "model.layers.{}.self_attn" ln_f_module: "model.norm" lm_head_module: "lm_head" mom2_dataset: "wikipedia" mom2_n_samples: 1000 mom2_dtype: "float16" model_parallel: false

[XeeKee]

In our local tests, this issue does indeed occur, and the speed of MEMIT suddenly slows down. I believe there is room for optimization here, but this is not a bug. I will try some optimization algorithms soon to speed it up, but it may take some time to optimize.

[zxlzr]

Hi, do you have any further issues?