fix: use W_org matrix only where needed (#398)

* fix: minor change

use `W_org` matrix where needed...

* Update model.py

* Update model.py

* fix: Windows hash, remove BOM marker

* docs: Add info about test cases

* feat: Tests for row_normalization PRE & NONE

* feat: CI hash files for row_normalization PRE & NONE models

* feat: Documentation instructions about test suite

* add recommendation

src/heretic/model.py

@@ -555,9 +555,11 @@ class Model:
                     # Flatten weight matrix to (out_features, in_features).
                     W = W.view(W.shape[0], -1)
 
-                    if self.settings.row_normalization != RowNormalization.NONE:
+                    if self.settings.row_normalization == RowNormalization.FULL:
                         # Keep a reference to the original weight matrix so we can subtract it later.
                         W_org = W
+
+                    if self.settings.row_normalization != RowNormalization.NONE:
                         # Get the row norms.
                         W_row_norms = LA.vector_norm(W, dim=1, keepdim=True)
                         # Normalize the weight matrix along the rows.

tests/README.md

@@ -1,17 +1,90 @@
-Run the tests with
+# Test Suite Guide
 
-```sh
-uv run run_tests.py
-```
+Whenever we change any code-logic related to `src/heretic/model.py` or `config.toml` *(e.g. `row_normalization`, `full_normalization_lora_rank`, `winsorization_quantile`, etc)* which can affect a model's reproduciblity; Use these tests which are designed to verify that those changes does not affect reproducibility, unless they are meant to (like when we'll integrate ARA branch in future).
 
-To update the hashes after a logic change, run the tests, then execute
+## How to test
 
-```sh
-cd TEST_DIR/model
+1. Choose any model from [tiny-random](https://huggingface.co/tiny-random) org which provides tiny models useful for debugging. 
+
+**Example**: [tiny-random/minicpm5](https://huggingface.co/tiny-random/minicpm5).
+
+> [!NOTE]
+> It is highly recommended to use a model which does not have a `special_tokens_map.json` file in the repo.
+> Because those files are almost always wrong in `tiny-random/*` models compared to the original model.
+
+2. Clone that model repository using Git and generate the SHA256 hashes using `sha256sum`:
+
+**On Linux**:
+
+```bash
 sha256sum -b * > ../SHA256SUMS.LABEL
 ```
 
-where `LABEL` describes the type of system you are running the tests on.
-Since PyTorch does not guarantee exact cross-system reproducibility regardless of configuration,
-multiple valid hashes can be provided for each output file. The above update must be performed
-for each `TEST_DIR` and on each type of system.
+**On Windows**:
+
+```bash
+sha256sum * | Out-File -Encoding utf8NoBOM ../SHA256SUMS.LABEL
+```
+
+> [!TIP]
+> On windows, `sha256sum` is generally pre-installed by *Git for windows*.
+
+**Verify with**:
+
+```bash
+Get-Command sha256sum`
+```
+
+**Expected**:
+
+```bash
+CommandType     Name                                               Version    Source
+-----------     ----                                               -------    ------
+Application     sha256sum.exe                                      0.0.0.0    C:\Program Files\Git\usr\bin\sha256sum...
+```
+
+> [!NOTE]
+> You must use Windows Powershell `v7.X` not the core which is `v5.1`. This is required for `-Encoding utf8NoBOM` to work.
+>
+> See [Differences between Windows PowerShell 5.1 and PowerShell 7.x](https://learn.microsoft.com/en-us/powershell/scripting/whats-new/differences-from-windows-powershell?view=powershell-7.6) documentation.
+
+Where `LABEL` describes the type of system you are running the tests on.
+
+**Example**: 
+
+- `SHA256SUMS.windows` (For windows)
+- `SHA256SUMS.ci` (For GitHub CI)
+- `SHA256SUMS.linux` (For linux)
+
+3. Run the tests with:
+
+```bash
+uv run run_tests.py
+```
+
+The output hashes *should FAIL* against the `Valid hashes` in `SHA256SUMS` file of the test model you added. This is expected since Heretic changes the model. Without **Step 2**, the test model's folder will simply be ignored because it will not have a hash SUMS file to compare against.
+
+4. After that go to the output `TEST_MODEL_DIR/model` folder and re-generate the Actual hashes based on the system you are using.
+
+```bash
+cd TEST_MODEL_DIR/model
+sha256sum -b * > ../SHA256SUMS.LABEL # or use windows command.
+```
+
+5. Re-run the tests with:
+
+```bash
+uv run run_tests.py
+```
+
+This time the tests *should PASS* because we added the new hashes which are expected to be reproduced on the same system.
+
+6. After that push the `SHA256SUMS.LABEL` files and wait for GitHub CI actions to run those tests.
+
+Since PyTorch does not guarantee exact cross-system reproducibility regardless of configuration, multiple valid hashes can be provided for each output file. The above update must be performed for each `TEST_MODEL_DIR` and on each type of system.
+
+For this, copy the `Actual hash` value for *each mismatched unidentical* file into a `SHA256SUMS.ci` file.
+
+7. After that push the `SHA256SUMS.ci` files and wait for GitHub CI actions to re-run those tests.
+
+This time the tests *should* PASS because we added the new hashes which are expected to be reproduced on CI.

tests/gemma-4e/SHA256SUMS.windows

@@ -1,4 +1,4 @@
-b16d3228a775c549ba97af41233a54e9de8dd2b65250f78346661d18b936a8b5 *chat_template.jinja
+b16d3228a775c549ba97af41233a54e9de8dd2b65250f78346661d18b936a8b5 *chat_template.jinja
 0094ad598a8043f84d82ad5c886547bca1d1d7f302d82f1491f83d388e89acd4 *config.json
 1a019c5d688d54cf01318eab88cb4345dfa52135eb1d83c2f54125469eb88d5c *generation_config.json
 effe36925f85ecb1e29bba84501a456bb49df21e4047be8b7ea3f6f88181fb65 *model.safetensors

tests/gemma-4e/config.toml

@@ -1,3 +1,6 @@
+# This test case is for Hybrid-Edge models.
+# After any change related to it, this test should PASS.
+
 model = "tiny-random/gemma-4e"
 model_commit = "3a207ada2c2cd95e9671942e84cf47ea58f0f6af"
 

tests/minicpm5/SHA256SUMS.ci

@@ -0,0 +1,6 @@
+7451a05cf1e28a79d97d7c0bc951028c0b1915119bf9046acd06a0e3d931f47c *chat_template.jinja
+fe6fd41d9f2ce5d6486748cf0330b574f37bf7d4e915f7b39d1af1a185cac3c3 *config.json
+c4c2ef5ae4a4e2dd10655a3b99d801a8a50497286ddd042ba35bcfefc44ad349 *generation_config.json
+1535a9b7a91b2cb39ad280dbd9a940e2609a0b423d5b924df4d664e579912802 *model.safetensors
+ad92aaa8d3032c98a9158b8c5e8682bed10027ed6463e4fb1320fe5384210873 *tokenizer.json
+3ad32522c384dbe35192bb69de9befbf3f523e99d4bb3f95da757671d4c28281 *tokenizer_config.json

tests/minicpm5/SHA256SUMS.windows

@@ -0,0 +1,6 @@
+d8db3ff45c4c68a0ba9dee962ff1a0adde9a2be55e0895306f6bd2b2756f5adb *chat_template.jinja
+a9d6f64bb9d0c02b553119e475615153af625b5c2a16ccb8fb8b3c2cc348f465 *config.json
+0e7611a1e8fd0a06a139b0572b2c55b885ba9fb7db2022873c3508aebfb488aa *generation_config.json
+411d95f42d3e31aef41c28314c8f0431c980687a97904d32b4ef57c42199720f *model.safetensors
+ad92aaa8d3032c98a9158b8c5e8682bed10027ed6463e4fb1320fe5384210873 *tokenizer.json
+aa083f3da10340925734e876e41e235c459329294ecd35d7511ec5868c1f14e3 *tokenizer_config.json

tests/minicpm5/config.toml

@@ -0,0 +1,46 @@
+# This test case is for row_normalization="none".
+# After any change related to it, this test should PASS.
+
+model = "tiny-random/minicpm5"
+model_commit = "52270c5ae5dde31255029cd5958591db057bd377"
+
+seed = 12345
+print_debug_information = true
+
+batch_size = 2
+max_response_length = 10
+kl_divergence_target = 0
+n_trials = 2
+n_startup_trials = 1
+
+export_strategy = "merge"
+checkpoint_action = "restart"
+trial_index = 0
+model_action = "save"
+save_directory = "model"
+
+row_normalization = "none"
+
+[good_prompts]
+dataset = "mlabonne/harmless_alpaca"
+commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
+split = "train[:5]"
+column = "text"
+
+[bad_prompts]
+dataset = "mlabonne/harmful_behaviors"
+commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
+split = "train[:5]"
+column = "text"
+
+[good_evaluation_prompts]
+dataset = "mlabonne/harmless_alpaca"
+commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
+split = "test[:5]"
+column = "text"
+
+[bad_evaluation_prompts]
+dataset = "mlabonne/harmful_behaviors"
+commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
+split = "test[:5]"
+column = "text"

tests/mistral-3/SHA256SUMS.windows

@@ -1,4 +1,4 @@
-72f84af4ea36b82409c35e31b584361534305ef7c0d90fce20d0dc38a7efead8 *chat_template.jinja
+72f84af4ea36b82409c35e31b584361534305ef7c0d90fce20d0dc38a7efead8 *chat_template.jinja
 e4c5278b361c57621253c27a2c3db358e1580aec8a14be8e19d4420a224137cf *config.json
 8dde85c000ae807be907421465826c7c63a39f6acf6d04a5a84efaf116ed4ef7 *generation_config.json
 29aff97d5633dead9e1ccd29a2cc153b4b7431d22f63c8d6cf60bc6547681cc9 *model.safetensors

tests/mistral-3/config.toml

@@ -1,3 +1,6 @@
+# This test case is for Dense models.
+# After any change related to it, this test should PASS.
+
 model = "tiny-random/mistral-3"
 model_commit = "931aa2e5c9668fc3679e56aa44972fe18597d55d"
 

tests/qwen2.5/SHA256SUMS.ci

@@ -0,0 +1,6 @@
+cd8e9439f0570856fd70470bf8889ebd8b5d1107207f67a5efb46e342330527f *chat_template.jinja
+45134b857367fdcb97c0179199848c353fc28f8b95ac2244ac8f45cca448d864 *config.json
+e81e23e025c38e825dcf8375861e26a90e804276e4db9ee390122a4fdc95dae7 *generation_config.json
+bd86541d817978c896bd3579e69ae6d41b6382eaf1646accf83d6feb16acb703 *model.safetensors
+f7f96da3a872b5e901575b2067c744ad336c3a3d77a21584d20024557b1bd7f0 *tokenizer.json
+04b1682c59acbd057f4c9072297faa73d56fc9de053094c659cdb4c464f58f86 *tokenizer_config.json

tests/qwen2.5/SHA256SUMS.windows

@@ -0,0 +1,6 @@
+8aa40ce145adb73cb3a75194dc0224702a95850ec5275cabb728496bbd749fc6 *chat_template.jinja
+e8f2fcd2681eb92233c0902866441f79a207b235f0b03364d41ebf8c53df62a0 *config.json
+3fec6d7004e5ae311864de130b62e32dac87569874c91b3fe9c46e9309345c1c *generation_config.json
+bd86541d817978c896bd3579e69ae6d41b6382eaf1646accf83d6feb16acb703 *model.safetensors
+f7f96da3a872b5e901575b2067c744ad336c3a3d77a21584d20024557b1bd7f0 *tokenizer.json
+154e5ff1e7c152d964edf30da854ea62465c767719ac8e97e58babf2d4fa9079 *tokenizer_config.json

tests/qwen2.5/config.toml

@@ -0,0 +1,46 @@
+# This test case is for row_normalization="pre".
+# After any change related to it, this test should PASS.
+
+model = "tiny-random/qwen2.5"
+model_commit = "7a6a3128ee4137a248d6d1582824592b87a81647"
+
+seed = 12345
+print_debug_information = true
+
+batch_size = 2
+max_response_length = 10
+kl_divergence_target = 0
+n_trials = 2
+n_startup_trials = 1
+
+export_strategy = "merge"
+checkpoint_action = "restart"
+trial_index = 0
+model_action = "save"
+save_directory = "model"
+
+row_normalization = "pre"
+
+[good_prompts]
+dataset = "mlabonne/harmless_alpaca"
+commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
+split = "train[:5]"
+column = "text"
+
+[bad_prompts]
+dataset = "mlabonne/harmful_behaviors"
+commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
+split = "train[:5]"
+column = "text"
+
+[good_evaluation_prompts]
+dataset = "mlabonne/harmless_alpaca"
+commit = "02c6a92cfcf11bb0c387334f8146d149d65b587f"
+split = "test[:5]"
+column = "text"
+
+[bad_evaluation_prompts]
+dataset = "mlabonne/harmful_behaviors"
+commit = "01cead01398926d81f7c52bdb790ee8cf77ebba7"
+split = "test[:5]"
+column = "text"

tests/qwen3.5-moe/SHA256SUMS.windows

@@ -1,4 +1,4 @@
-a92e1dd97cb1cb175c9b70c0828e146bea4371c2643319b661b777e89811972e *chat_template.jinja
+a92e1dd97cb1cb175c9b70c0828e146bea4371c2643319b661b777e89811972e *chat_template.jinja
 b75e911805663da79fb9fbbbcc917b8f1a285d2da54d95c2c63ea7c1ffe9a05a *config.json
 2cbd9df0e99570efcced23b8d777bdf1fc692efda54b21eb59ad56ade76c9db6 *generation_config.json
 5f099b32807d0b84ed90765ca0ed53f8771da4738767bc1940486fec954570cf *model.safetensors

tests/qwen3.5-moe/config.toml

@@ -1,3 +1,6 @@
+# This test case is for MoE models.
+# After any change related to it, this test should PASS.
+
 model = "tiny-random/qwen3.5-moe"
 model_commit = "2ebfa8d9717238c5dda927008104fa172a149050"