Fix norms and randomize tests for subnet extraction

whittle-org / whittle

Python library to compress LitGPT models for resource efficient inference.

https://whittle-org.github.io/whittle/latest/

Apache License 2.0

10 stars 4 forks source link

Fix norms and randomize tests for subnet extraction #150

Closed rheasukthanker closed 3 weeks ago

rheasukthanker commented 3 weeks ago

Describe the bug In /whittle/models/gpt/extract.py weights and biases or norm layers are not handled properly. The tests still pass because they initialize the model with ones.

Possible solution

Support rmsnorm and layernorm layers
Randomize weights in extract subnetwork tests