diff --git a/src/heretic/model.py b/src/heretic/model.py index 0d0644c..3d7c026 100644 --- a/src/heretic/model.py +++ b/src/heretic/model.py @@ -437,6 +437,10 @@ class Model: with suppress(Exception): try_add("attn.o_proj", layer.linear_attn.out_proj) # ty: ignore[unresolved-attribute] + # DeepSeek V4. + with suppress(Exception): + try_add("attn.o_proj", layer.self_attn.o_b_proj) # ty: ignore[unresolved-attribute] + # Most dense models. with suppress(Exception): try_add("mlp.down_proj", layer.mlp.down_proj) # ty: ignore[unresolved-attribute] @@ -446,6 +450,10 @@ class Model: for expert in layer.mlp.experts: # ty:ignore[not-iterable, unresolved-attribute] try_add("mlp.down_proj", expert.down_proj) # ty: ignore[unresolved-attribute] + # Shared experts of MoE models (e.g. DeepSeek V4 and GLM 5.3). + with suppress(Exception): + try_add("mlp.down_proj", layer.mlp.shared_experts.down_proj) # ty: ignore[unresolved-attribute] + # Phi-3.5-MoE (and possibly others). with suppress(Exception): for expert in layer.block_sparse_moe.experts: # ty:ignore[not-iterable, unresolved-attribute] @@ -632,6 +640,20 @@ class Model: # This cast is valid because we passed output_hidden_states=True above. hidden_states = cast(tuple[tuple[FloatTensor]], outputs.hidden_states)[0] + # Models with Hyper-Connections (e.g. DeepSeek V4 and GLM 5.3) have several + # residual streams, so their hidden states have shape + # (prompt, position, stream, component). We use the mean of the streams. + # The last entry has no stream dimension, because Transformers replaces it + # with the model output, where the streams have already been merged. + hidden_states = tuple( + ( + layer_hidden_states.mean(dim=2) + if layer_hidden_states.dim() == 4 + else layer_hidden_states + ) + for layer_hidden_states in hidden_states + ) + # The returned tensor has shape (prompt, layer, component). residuals = torch.stack( # layer_hidden_states has shape (prompt, position, component), diff --git a/tests/deepseek-v4-ara/SHA256SUMS.ci b/tests/deepseek-v4-ara/SHA256SUMS.ci new file mode 100644 index 0000000..9f6187a --- /dev/null +++ b/tests/deepseek-v4-ara/SHA256SUMS.ci @@ -0,0 +1 @@ +03ca3c7b8ad59ddae470c86849f6623e08ccf7ef911ab239fe339e3fbedbb8b9 *adapter_model.safetensors diff --git a/tests/deepseek-v4-ara/SHA256SUMS.ci2 b/tests/deepseek-v4-ara/SHA256SUMS.ci2 new file mode 100644 index 0000000..46dabeb --- /dev/null +++ b/tests/deepseek-v4-ara/SHA256SUMS.ci2 @@ -0,0 +1 @@ +f01ac38a0efdfd876f604f3ecba4f7239c1e8b4c199d0711291dfba52eb107a4 *adapter_model.safetensors diff --git a/tests/deepseek-v4-ara/SHA256SUMS.ci3 b/tests/deepseek-v4-ara/SHA256SUMS.ci3 new file mode 100644 index 0000000..b5b6aaa --- /dev/null +++ b/tests/deepseek-v4-ara/SHA256SUMS.ci3 @@ -0,0 +1 @@ +a4e17011af1957d0e66505b302dc4b6f098e85a823a95ef1b9d93e5c3623b9c7 *adapter_model.safetensors diff --git a/tests/deepseek-v4-ara/SHA256SUMS.linux b/tests/deepseek-v4-ara/SHA256SUMS.linux new file mode 100644 index 0000000..bb242ff --- /dev/null +++ b/tests/deepseek-v4-ara/SHA256SUMS.linux @@ -0,0 +1 @@ +de63a81fb63f779edf88631fb0a8ae28badc1b07df9d49a7151056fc488d4b7a *adapter_model.safetensors diff --git a/tests/deepseek-v4-ara/config.toml b/tests/deepseek-v4-ara/config.toml new file mode 100644 index 0000000..35432c1 --- /dev/null +++ b/tests/deepseek-v4-ara/config.toml @@ -0,0 +1,59 @@ +# This test case is for models with Hyper-Connections (multiple residual streams), +# using the ARA modifier, which handles residuals differently than Abliteration. +# After any change related to it, this test should PASS. + +model = "tiny-random/deepseek-v4" +model_commit = "d20cabaa1c8a9f235e9bb50f1d5e98eeb6a530d5" + +seed = 12345 +print_debug_information = true + +batch_size = 2 +max_response_length = 10 + +n_trials = 2 +n_startup_trials = 1 + +# Transformers can't yet save DeepSeek V4 after dequantizing its FP8 weights, +# so this test checks the LoRA adapter instead of the merged model. +export_strategy = "adapter" +checkpoint_action = "restart" +trial_index = 0 +model_action = "save" +save_directory = "model" + +[[response_prefix_test_prompts]] +dataset = "mlabonne/harmless_alpaca" +commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f" +split = "train[:5]" +column = "text" + +[[response_prefix_test_prompts]] +dataset = "mlabonne/harmful_behaviors" +commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7" +split = "train[:5]" +column = "text" + +[scorer.KLDivergence.prompts] +dataset = "mlabonne/harmless_alpaca" +commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f" +split = "test[:5]" +column = "text" + +[scorer.KeywordRate.prompts] +dataset = "mlabonne/harmful_behaviors" +commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7" +split = "test[:5]" +column = "text" + +[modifier.ARA.good_prompts] +dataset = "mlabonne/harmless_alpaca" +commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f" +split = "train[:5]" +column = "text" + +[modifier.ARA.bad_prompts] +dataset = "mlabonne/harmful_behaviors" +commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7" +split = "train[:5]" +column = "text" diff --git a/tests/deepseek-v4/SHA256SUMS.ci b/tests/deepseek-v4/SHA256SUMS.ci new file mode 100644 index 0000000..66c73e1 --- /dev/null +++ b/tests/deepseek-v4/SHA256SUMS.ci @@ -0,0 +1 @@ +221e865de7fbd712de6d0f19012e3da915b326eb797a2b430afa146c6d1d6cf2 *adapter_model.safetensors diff --git a/tests/deepseek-v4/SHA256SUMS.ci2 b/tests/deepseek-v4/SHA256SUMS.ci2 new file mode 100644 index 0000000..375f387 --- /dev/null +++ b/tests/deepseek-v4/SHA256SUMS.ci2 @@ -0,0 +1 @@ +e73ed117efc71ba636195891caa393abee29eb136273dcf2a48ebec0c3611717 *adapter_model.safetensors diff --git a/tests/deepseek-v4/SHA256SUMS.ci3 b/tests/deepseek-v4/SHA256SUMS.ci3 new file mode 100644 index 0000000..8a3c9a1 --- /dev/null +++ b/tests/deepseek-v4/SHA256SUMS.ci3 @@ -0,0 +1 @@ +e62c18d8cd28576234d45978ea273ce446be0eebc74b170080a8f9238ec2a233 *adapter_model.safetensors diff --git a/tests/deepseek-v4/SHA256SUMS.ci4 b/tests/deepseek-v4/SHA256SUMS.ci4 new file mode 100644 index 0000000..63592ff --- /dev/null +++ b/tests/deepseek-v4/SHA256SUMS.ci4 @@ -0,0 +1 @@ +6b923a6e0d39ed7664d36a913c7e885c684d8996a1958fc4d470afdcc586d70f *adapter_model.safetensors diff --git a/tests/deepseek-v4/SHA256SUMS.linux b/tests/deepseek-v4/SHA256SUMS.linux new file mode 100644 index 0000000..db82abf --- /dev/null +++ b/tests/deepseek-v4/SHA256SUMS.linux @@ -0,0 +1 @@ +a03e314411ba09c75f1ee62da54d260ac8b68bd5ddedaa0f709c4bb975b06143 *adapter_model.safetensors diff --git a/tests/deepseek-v4/SHA256SUMS.windows b/tests/deepseek-v4/SHA256SUMS.windows new file mode 100644 index 0000000..2389aa2 --- /dev/null +++ b/tests/deepseek-v4/SHA256SUMS.windows @@ -0,0 +1 @@ +9de2b756fa643074ac3641d97076d39c9f89e08a24a39a546642a59d975442f5 *adapter_model.safetensors diff --git a/tests/deepseek-v4/config.toml b/tests/deepseek-v4/config.toml new file mode 100644 index 0000000..40969e1 --- /dev/null +++ b/tests/deepseek-v4/config.toml @@ -0,0 +1,62 @@ +# This test case is for models with Hyper-Connections (multiple residual streams). +# After any change related to it, this test should PASS. + +model = "tiny-random/deepseek-v4" +model_commit = "d20cabaa1c8a9f235e9bb50f1d5e98eeb6a530d5" + +seed = 12345 +print_debug_information = true + +batch_size = 2 +max_response_length = 10 + +modifiers = [ + { plugin = "heretic.modifiers.abliteration.Abliteration" }, +] + +n_trials = 2 +n_startup_trials = 1 + +# Transformers can't yet save DeepSeek V4 after dequantizing its FP8 weights, +# so this test checks the LoRA adapter instead of the merged model. +export_strategy = "adapter" +checkpoint_action = "restart" +trial_index = 0 +model_action = "save" +save_directory = "model" + +[[response_prefix_test_prompts]] +dataset = "mlabonne/harmless_alpaca" +commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f" +split = "train[:5]" +column = "text" + +[[response_prefix_test_prompts]] +dataset = "mlabonne/harmful_behaviors" +commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7" +split = "train[:5]" +column = "text" + +[scorer.KLDivergence.prompts] +dataset = "mlabonne/harmless_alpaca" +commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f" +split = "test[:5]" +column = "text" + +[scorer.KeywordRate.prompts] +dataset = "mlabonne/harmful_behaviors" +commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7" +split = "test[:5]" +column = "text" + +[modifier.Abliteration.good_prompts] +dataset = "mlabonne/harmless_alpaca" +commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f" +split = "train[:5]" +column = "text" + +[modifier.Abliteration.bad_prompts] +dataset = "mlabonne/harmful_behaviors" +commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7" +split = "train[:5]" +column = "text"