Defuser 0.0.17__tar.gz → 0.0.19__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {defuser-0.0.17 → defuser-0.0.19}/Defuser.egg-info/PKG-INFO +10 -5
- {defuser-0.0.17 → defuser-0.0.19}/Defuser.egg-info/SOURCES.txt +2 -0
- {defuser-0.0.17 → defuser-0.0.19}/PKG-INFO +10 -5
- {defuser-0.0.17 → defuser-0.0.19}/README.md +8 -4
- {defuser-0.0.17 → defuser-0.0.19}/defuser/__init__.py +13 -3
- {defuser-0.0.17 → defuser-0.0.19}/defuser/checkpoint_ops.py +4 -0
- {defuser-0.0.17 → defuser-0.0.19}/defuser/defuser.py +21 -3
- {defuser-0.0.17 → defuser-0.0.19}/defuser/model_registry.py +6 -0
- defuser-0.0.19/defuser/modeling/__init__.py +4 -0
- {defuser-0.0.17 → defuser-0.0.19}/defuser/modeling/glm4v.py +4 -0
- {defuser-0.0.17 → defuser-0.0.19}/defuser/modeling/moe_experts_interface.py +60 -29
- defuser-0.0.19/defuser/modeling/unfused_moe/__init__.py +4 -0
- defuser-0.0.19/defuser/modeling/unfused_moe/glm4_moe_lite.py +100 -0
- defuser-0.0.19/defuser/utils/__init__.py +4 -0
- {defuser-0.0.17 → defuser-0.0.19}/pyproject.toml +2 -1
- defuser-0.0.19/setup.py +7 -0
- {defuser-0.0.17 → defuser-0.0.19}/tests/test_candidate_coverage.py +31 -3
- {defuser-0.0.17 → defuser-0.0.19}/tests/test_convert_model.py +32 -2
- {defuser-0.0.17 → defuser-0.0.19}/tests/test_meta_model_defusion.py +23 -4
- defuser-0.0.19/tests/test_moe_experts_interface.py +62 -0
- defuser-0.0.17/defuser/modeling/__init__.py +0 -0
- defuser-0.0.17/defuser/modeling/unfused_moe/__init__.py +0 -0
- defuser-0.0.17/defuser/utils/__init__.py +0 -0
- defuser-0.0.17/setup.py +0 -3
- {defuser-0.0.17 → defuser-0.0.19}/Defuser.egg-info/dependency_links.txt +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/Defuser.egg-info/requires.txt +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/Defuser.egg-info/top_level.txt +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/LICENSE +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/defuser/modeling/model_patches.py +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/defuser/modeling/replace_modules.py +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/defuser/modeling/runtime_defusion.py +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/defuser/modeling/unfused_moe/common.py +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/defuser/modeling/unfused_moe/glm4_moe.py +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/defuser/modeling/unfused_moe/mixtral.py +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/defuser/modeling/unfused_moe/qwen2_moe.py +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/defuser/modeling/unfused_moe/qwen3_moe.py +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/defuser/modeling/unfused_moe/qwen3_next.py +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/defuser/modeling/unfused_moe/qwen3_omni_moe.py +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/defuser/modeling/update_module.py +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/defuser/utils/common.py +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/defuser/utils/device.py +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/defuser/utils/hf.py +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/defuser/utils/model.py +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/setup.cfg +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/tests/test_device_utils.py +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/tests/test_filter_rules.py +0 -0
- {defuser-0.0.17 → defuser-0.0.19}/tests/test_replace_modules_tracker.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: Defuser
|
|
3
|
-
Version: 0.0.
|
|
3
|
+
Version: 0.0.19
|
|
4
4
|
Summary: Model defuser helper for HF Transformers.
|
|
5
5
|
Author-email: ModelCloud <qubitium@modelcloud.ai>
|
|
6
6
|
License-Expression: Apache-2.0
|
|
@@ -13,6 +13,7 @@ Classifier: Programming Language :: Python :: 3.11
|
|
|
13
13
|
Classifier: Programming Language :: Python :: 3.12
|
|
14
14
|
Classifier: Programming Language :: Python :: 3.13
|
|
15
15
|
Classifier: Programming Language :: Python :: 3.14
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.15
|
|
16
17
|
Classifier: Programming Language :: C
|
|
17
18
|
Classifier: Topic :: Software Development :: Libraries
|
|
18
19
|
Classifier: Topic :: Text Processing :: Linguistic
|
|
@@ -48,7 +49,7 @@ Depending on the model family, Defuser can:
|
|
|
48
49
|
|
|
49
50
|
- patch a supported model class before load so HF instantiates a defused block directly
|
|
50
51
|
- split fused tensors such as `gate_up_proj` into `gate_proj` + `up_proj`
|
|
51
|
-
- convert 3D expert tensors into numbered expert `nn.Linear` modules
|
|
52
|
+
- convert 3D expert tensors, including registered expert buffers, into numbered expert `nn.Linear` modules
|
|
52
53
|
- preserve the original fused math while presenting a naive module structure again
|
|
53
54
|
|
|
54
55
|
Public API:
|
|
@@ -58,8 +59,9 @@ from defuser import convert_model, replace_fused_blocks
|
|
|
58
59
|
```
|
|
59
60
|
|
|
60
61
|
- `replace_fused_blocks(model_type)` patches supported HF model classes before `from_pretrained()` or direct model construction.
|
|
61
|
-
- `convert_model(model, cleanup_original=
|
|
62
|
+
- `convert_model(model, cleanup_original=False, max_layers=None, filter=None)` converts an already loaded model in place. This is the runtime defusion path for supported post-load expert and MLP conversions, including `qwen3_5_moe` style checkpoints.
|
|
62
63
|
- Defuser is designed and CI-tested for `transformers>=5.3.0`, and support is only offered for that version range. Older versions log a warning on these public APIs and are skipped as unsupported.
|
|
64
|
+
- Some model families appear in both support tables. Full models can be prepatched with `replace_fused_blocks(...)`, while standalone fused expert modules from those same families can still be runtime-defused with `convert_model(...)`.
|
|
63
65
|
|
|
64
66
|
`filter` is an optional list of PCRE regex rules evaluated against full module paths such as `model.layers.0.mlp.experts`:
|
|
65
67
|
|
|
@@ -71,13 +73,14 @@ from defuser import convert_model, replace_fused_blocks
|
|
|
71
73
|
|
|
72
74
|
## Supported Models
|
|
73
75
|
|
|
74
|
-
Defuser currently supports the following `transformers
|
|
76
|
+
Defuser currently supports the following `transformers>=5.3.0` `model_type` values.
|
|
75
77
|
|
|
76
78
|
### `replace_fused_blocks(model_type)` before load
|
|
77
79
|
|
|
78
80
|
| Model type | Defused op performed |
|
|
79
81
|
| --- | --- |
|
|
80
82
|
| `glm4_moe` | Replaces `Glm4MoeMoE` with a defused per-expert linear MoE block. |
|
|
83
|
+
| `glm4_moe_lite` | Replaces `Glm4MoeLiteMoE` with a defused per-expert linear MoE block.|
|
|
81
84
|
| `glm4v` | Replaces the fused text MLP with split `gate_proj`, `up_proj`, and `down_proj` layers. Also splits fused checkpoint `mlp.gate_up_proj.weight` into `mlp.gate_proj.weight` + `mlp.up_proj.weight`. |
|
|
82
85
|
| `mixtral` | Replaces `MixtralSparseMoeBlock` with `LinearMixtralSparseMoeBlock`. Also remaps legacy Mixtral checkpoint keys and splits fused expert `gate_up_proj` tensors into per-expert `gate_proj` and `up_proj`, plus per-expert `down_proj`. |
|
|
83
86
|
| `qwen2_moe` | Replaces `Qwen2MoeSparseMoeBlock` with a defused per-expert linear MoE block. |
|
|
@@ -89,7 +92,7 @@ Defuser currently supports the following `transformers==5.3.0` `model_type` valu
|
|
|
89
92
|
|
|
90
93
|
| Pattern | Supported model types | Defused op performed |
|
|
91
94
|
| --- | --- | --- |
|
|
92
|
-
| Standard routed expert tensors | `deepseek_v2`, `dots1`, `ernie4_5_moe`, `ernie4_5_vl_moe`, `exaone_moe`, `flex_olmo`, `glm4_moe_lite`, `glm4v_moe`, `hunyuan_v1_moe`, `jamba`, `lfm2_moe`, `minimax`, `minimax_m2`, `olmoe`, `qwen3_vl_moe`, `solar_open` | Splits fused expert tensors into numbered expert `nn.Linear` modules with per-expert `gate_proj`, `up_proj`, and `down_proj`. |
|
|
95
|
+
| Standard routed expert tensors | `deepseek_v2`, `dots1`, `ernie4_5_moe`, `ernie4_5_vl_moe`, `exaone_moe`, `flex_olmo`, `glm4_moe_lite`, `glm4v_moe`, `hunyuan_v1_moe`, `jamba`, `lfm2_moe`, `minimax`, `minimax_m2`, `olmoe`, `qwen3_vl_moe`, `solar_open` | Splits fused expert tensors or registered expert buffers into numbered expert `nn.Linear` modules with per-expert `gate_proj`, `up_proj`, and `down_proj`. |
|
|
93
96
|
| Mixed sparse and shared experts | `deepseek_v3`, `glm_moe_dsa`, `qwen3_5_moe`, `qwen3_5_moe_text` | Runtime expert tensor defusion for routed experts while preserving the model's shared-expert path. |
|
|
94
97
|
| Transposed or packed expert tensors | `gpt_oss`, `phimoe` | Splits transposed fused expert `gate_up_proj` tensors into per-expert `gate_proj` + `up_proj`, preserves expert bias when present, and converts expert tensors into numbered expert `nn.Linear` modules. |
|
|
95
98
|
| Flattened expert layout | `dbrx` | Rebuilds the flattened DBRX expert FFN weights into numbered expert `gate_proj`, `up_proj`, and `down_proj` `nn.Linear` modules. |
|
|
@@ -124,6 +127,8 @@ converted = convert_model(model)
|
|
|
124
127
|
print(converted) # True when runtime defusion happened
|
|
125
128
|
```
|
|
126
129
|
|
|
130
|
+
`convert_model(model)` also preserves meta-device construction for supported meta-initialized models, so structural validation can run without materializing weights.
|
|
131
|
+
|
|
127
132
|
Use `filter` when only specific blocks should be defused:
|
|
128
133
|
|
|
129
134
|
```python
|
|
@@ -21,6 +21,7 @@ defuser/modeling/update_module.py
|
|
|
21
21
|
defuser/modeling/unfused_moe/__init__.py
|
|
22
22
|
defuser/modeling/unfused_moe/common.py
|
|
23
23
|
defuser/modeling/unfused_moe/glm4_moe.py
|
|
24
|
+
defuser/modeling/unfused_moe/glm4_moe_lite.py
|
|
24
25
|
defuser/modeling/unfused_moe/mixtral.py
|
|
25
26
|
defuser/modeling/unfused_moe/qwen2_moe.py
|
|
26
27
|
defuser/modeling/unfused_moe/qwen3_moe.py
|
|
@@ -36,4 +37,5 @@ tests/test_convert_model.py
|
|
|
36
37
|
tests/test_device_utils.py
|
|
37
38
|
tests/test_filter_rules.py
|
|
38
39
|
tests/test_meta_model_defusion.py
|
|
40
|
+
tests/test_moe_experts_interface.py
|
|
39
41
|
tests/test_replace_modules_tracker.py
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: Defuser
|
|
3
|
-
Version: 0.0.
|
|
3
|
+
Version: 0.0.19
|
|
4
4
|
Summary: Model defuser helper for HF Transformers.
|
|
5
5
|
Author-email: ModelCloud <qubitium@modelcloud.ai>
|
|
6
6
|
License-Expression: Apache-2.0
|
|
@@ -13,6 +13,7 @@ Classifier: Programming Language :: Python :: 3.11
|
|
|
13
13
|
Classifier: Programming Language :: Python :: 3.12
|
|
14
14
|
Classifier: Programming Language :: Python :: 3.13
|
|
15
15
|
Classifier: Programming Language :: Python :: 3.14
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.15
|
|
16
17
|
Classifier: Programming Language :: C
|
|
17
18
|
Classifier: Topic :: Software Development :: Libraries
|
|
18
19
|
Classifier: Topic :: Text Processing :: Linguistic
|
|
@@ -48,7 +49,7 @@ Depending on the model family, Defuser can:
|
|
|
48
49
|
|
|
49
50
|
- patch a supported model class before load so HF instantiates a defused block directly
|
|
50
51
|
- split fused tensors such as `gate_up_proj` into `gate_proj` + `up_proj`
|
|
51
|
-
- convert 3D expert tensors into numbered expert `nn.Linear` modules
|
|
52
|
+
- convert 3D expert tensors, including registered expert buffers, into numbered expert `nn.Linear` modules
|
|
52
53
|
- preserve the original fused math while presenting a naive module structure again
|
|
53
54
|
|
|
54
55
|
Public API:
|
|
@@ -58,8 +59,9 @@ from defuser import convert_model, replace_fused_blocks
|
|
|
58
59
|
```
|
|
59
60
|
|
|
60
61
|
- `replace_fused_blocks(model_type)` patches supported HF model classes before `from_pretrained()` or direct model construction.
|
|
61
|
-
- `convert_model(model, cleanup_original=
|
|
62
|
+
- `convert_model(model, cleanup_original=False, max_layers=None, filter=None)` converts an already loaded model in place. This is the runtime defusion path for supported post-load expert and MLP conversions, including `qwen3_5_moe` style checkpoints.
|
|
62
63
|
- Defuser is designed and CI-tested for `transformers>=5.3.0`, and support is only offered for that version range. Older versions log a warning on these public APIs and are skipped as unsupported.
|
|
64
|
+
- Some model families appear in both support tables. Full models can be prepatched with `replace_fused_blocks(...)`, while standalone fused expert modules from those same families can still be runtime-defused with `convert_model(...)`.
|
|
63
65
|
|
|
64
66
|
`filter` is an optional list of PCRE regex rules evaluated against full module paths such as `model.layers.0.mlp.experts`:
|
|
65
67
|
|
|
@@ -71,13 +73,14 @@ from defuser import convert_model, replace_fused_blocks
|
|
|
71
73
|
|
|
72
74
|
## Supported Models
|
|
73
75
|
|
|
74
|
-
Defuser currently supports the following `transformers
|
|
76
|
+
Defuser currently supports the following `transformers>=5.3.0` `model_type` values.
|
|
75
77
|
|
|
76
78
|
### `replace_fused_blocks(model_type)` before load
|
|
77
79
|
|
|
78
80
|
| Model type | Defused op performed |
|
|
79
81
|
| --- | --- |
|
|
80
82
|
| `glm4_moe` | Replaces `Glm4MoeMoE` with a defused per-expert linear MoE block. |
|
|
83
|
+
| `glm4_moe_lite` | Replaces `Glm4MoeLiteMoE` with a defused per-expert linear MoE block.|
|
|
81
84
|
| `glm4v` | Replaces the fused text MLP with split `gate_proj`, `up_proj`, and `down_proj` layers. Also splits fused checkpoint `mlp.gate_up_proj.weight` into `mlp.gate_proj.weight` + `mlp.up_proj.weight`. |
|
|
82
85
|
| `mixtral` | Replaces `MixtralSparseMoeBlock` with `LinearMixtralSparseMoeBlock`. Also remaps legacy Mixtral checkpoint keys and splits fused expert `gate_up_proj` tensors into per-expert `gate_proj` and `up_proj`, plus per-expert `down_proj`. |
|
|
83
86
|
| `qwen2_moe` | Replaces `Qwen2MoeSparseMoeBlock` with a defused per-expert linear MoE block. |
|
|
@@ -89,7 +92,7 @@ Defuser currently supports the following `transformers==5.3.0` `model_type` valu
|
|
|
89
92
|
|
|
90
93
|
| Pattern | Supported model types | Defused op performed |
|
|
91
94
|
| --- | --- | --- |
|
|
92
|
-
| Standard routed expert tensors | `deepseek_v2`, `dots1`, `ernie4_5_moe`, `ernie4_5_vl_moe`, `exaone_moe`, `flex_olmo`, `glm4_moe_lite`, `glm4v_moe`, `hunyuan_v1_moe`, `jamba`, `lfm2_moe`, `minimax`, `minimax_m2`, `olmoe`, `qwen3_vl_moe`, `solar_open` | Splits fused expert tensors into numbered expert `nn.Linear` modules with per-expert `gate_proj`, `up_proj`, and `down_proj`. |
|
|
95
|
+
| Standard routed expert tensors | `deepseek_v2`, `dots1`, `ernie4_5_moe`, `ernie4_5_vl_moe`, `exaone_moe`, `flex_olmo`, `glm4_moe_lite`, `glm4v_moe`, `hunyuan_v1_moe`, `jamba`, `lfm2_moe`, `minimax`, `minimax_m2`, `olmoe`, `qwen3_vl_moe`, `solar_open` | Splits fused expert tensors or registered expert buffers into numbered expert `nn.Linear` modules with per-expert `gate_proj`, `up_proj`, and `down_proj`. |
|
|
93
96
|
| Mixed sparse and shared experts | `deepseek_v3`, `glm_moe_dsa`, `qwen3_5_moe`, `qwen3_5_moe_text` | Runtime expert tensor defusion for routed experts while preserving the model's shared-expert path. |
|
|
94
97
|
| Transposed or packed expert tensors | `gpt_oss`, `phimoe` | Splits transposed fused expert `gate_up_proj` tensors into per-expert `gate_proj` + `up_proj`, preserves expert bias when present, and converts expert tensors into numbered expert `nn.Linear` modules. |
|
|
95
98
|
| Flattened expert layout | `dbrx` | Rebuilds the flattened DBRX expert FFN weights into numbered expert `gate_proj`, `up_proj`, and `down_proj` `nn.Linear` modules. |
|
|
@@ -124,6 +127,8 @@ converted = convert_model(model)
|
|
|
124
127
|
print(converted) # True when runtime defusion happened
|
|
125
128
|
```
|
|
126
129
|
|
|
130
|
+
`convert_model(model)` also preserves meta-device construction for supported meta-initialized models, so structural validation can run without materializing weights.
|
|
131
|
+
|
|
127
132
|
Use `filter` when only specific blocks should be defused:
|
|
128
133
|
|
|
129
134
|
```python
|
|
@@ -23,7 +23,7 @@ Depending on the model family, Defuser can:
|
|
|
23
23
|
|
|
24
24
|
- patch a supported model class before load so HF instantiates a defused block directly
|
|
25
25
|
- split fused tensors such as `gate_up_proj` into `gate_proj` + `up_proj`
|
|
26
|
-
- convert 3D expert tensors into numbered expert `nn.Linear` modules
|
|
26
|
+
- convert 3D expert tensors, including registered expert buffers, into numbered expert `nn.Linear` modules
|
|
27
27
|
- preserve the original fused math while presenting a naive module structure again
|
|
28
28
|
|
|
29
29
|
Public API:
|
|
@@ -33,8 +33,9 @@ from defuser import convert_model, replace_fused_blocks
|
|
|
33
33
|
```
|
|
34
34
|
|
|
35
35
|
- `replace_fused_blocks(model_type)` patches supported HF model classes before `from_pretrained()` or direct model construction.
|
|
36
|
-
- `convert_model(model, cleanup_original=
|
|
36
|
+
- `convert_model(model, cleanup_original=False, max_layers=None, filter=None)` converts an already loaded model in place. This is the runtime defusion path for supported post-load expert and MLP conversions, including `qwen3_5_moe` style checkpoints.
|
|
37
37
|
- Defuser is designed and CI-tested for `transformers>=5.3.0`, and support is only offered for that version range. Older versions log a warning on these public APIs and are skipped as unsupported.
|
|
38
|
+
- Some model families appear in both support tables. Full models can be prepatched with `replace_fused_blocks(...)`, while standalone fused expert modules from those same families can still be runtime-defused with `convert_model(...)`.
|
|
38
39
|
|
|
39
40
|
`filter` is an optional list of PCRE regex rules evaluated against full module paths such as `model.layers.0.mlp.experts`:
|
|
40
41
|
|
|
@@ -46,13 +47,14 @@ from defuser import convert_model, replace_fused_blocks
|
|
|
46
47
|
|
|
47
48
|
## Supported Models
|
|
48
49
|
|
|
49
|
-
Defuser currently supports the following `transformers
|
|
50
|
+
Defuser currently supports the following `transformers>=5.3.0` `model_type` values.
|
|
50
51
|
|
|
51
52
|
### `replace_fused_blocks(model_type)` before load
|
|
52
53
|
|
|
53
54
|
| Model type | Defused op performed |
|
|
54
55
|
| --- | --- |
|
|
55
56
|
| `glm4_moe` | Replaces `Glm4MoeMoE` with a defused per-expert linear MoE block. |
|
|
57
|
+
| `glm4_moe_lite` | Replaces `Glm4MoeLiteMoE` with a defused per-expert linear MoE block.|
|
|
56
58
|
| `glm4v` | Replaces the fused text MLP with split `gate_proj`, `up_proj`, and `down_proj` layers. Also splits fused checkpoint `mlp.gate_up_proj.weight` into `mlp.gate_proj.weight` + `mlp.up_proj.weight`. |
|
|
57
59
|
| `mixtral` | Replaces `MixtralSparseMoeBlock` with `LinearMixtralSparseMoeBlock`. Also remaps legacy Mixtral checkpoint keys and splits fused expert `gate_up_proj` tensors into per-expert `gate_proj` and `up_proj`, plus per-expert `down_proj`. |
|
|
58
60
|
| `qwen2_moe` | Replaces `Qwen2MoeSparseMoeBlock` with a defused per-expert linear MoE block. |
|
|
@@ -64,7 +66,7 @@ Defuser currently supports the following `transformers==5.3.0` `model_type` valu
|
|
|
64
66
|
|
|
65
67
|
| Pattern | Supported model types | Defused op performed |
|
|
66
68
|
| --- | --- | --- |
|
|
67
|
-
| Standard routed expert tensors | `deepseek_v2`, `dots1`, `ernie4_5_moe`, `ernie4_5_vl_moe`, `exaone_moe`, `flex_olmo`, `glm4_moe_lite`, `glm4v_moe`, `hunyuan_v1_moe`, `jamba`, `lfm2_moe`, `minimax`, `minimax_m2`, `olmoe`, `qwen3_vl_moe`, `solar_open` | Splits fused expert tensors into numbered expert `nn.Linear` modules with per-expert `gate_proj`, `up_proj`, and `down_proj`. |
|
|
69
|
+
| Standard routed expert tensors | `deepseek_v2`, `dots1`, `ernie4_5_moe`, `ernie4_5_vl_moe`, `exaone_moe`, `flex_olmo`, `glm4_moe_lite`, `glm4v_moe`, `hunyuan_v1_moe`, `jamba`, `lfm2_moe`, `minimax`, `minimax_m2`, `olmoe`, `qwen3_vl_moe`, `solar_open` | Splits fused expert tensors or registered expert buffers into numbered expert `nn.Linear` modules with per-expert `gate_proj`, `up_proj`, and `down_proj`. |
|
|
68
70
|
| Mixed sparse and shared experts | `deepseek_v3`, `glm_moe_dsa`, `qwen3_5_moe`, `qwen3_5_moe_text` | Runtime expert tensor defusion for routed experts while preserving the model's shared-expert path. |
|
|
69
71
|
| Transposed or packed expert tensors | `gpt_oss`, `phimoe` | Splits transposed fused expert `gate_up_proj` tensors into per-expert `gate_proj` + `up_proj`, preserves expert bias when present, and converts expert tensors into numbered expert `nn.Linear` modules. |
|
|
70
72
|
| Flattened expert layout | `dbrx` | Rebuilds the flattened DBRX expert FFN weights into numbered expert `gate_proj`, `up_proj`, and `down_proj` `nn.Linear` modules. |
|
|
@@ -99,6 +101,8 @@ converted = convert_model(model)
|
|
|
99
101
|
print(converted) # True when runtime defusion happened
|
|
100
102
|
```
|
|
101
103
|
|
|
104
|
+
`convert_model(model)` also preserves meta-device construction for supported meta-initialized models, so structural validation can run without materializing weights.
|
|
105
|
+
|
|
102
106
|
Use `filter` when only specific blocks should be defused:
|
|
103
107
|
|
|
104
108
|
```python
|
|
@@ -3,19 +3,29 @@
|
|
|
3
3
|
# SPDX-License-Identifier: Apache-2.0
|
|
4
4
|
# Contact: qubitium@modelcloud.ai, x.com/qubitium
|
|
5
5
|
|
|
6
|
-
from defuser.utils.common import env_flag
|
|
6
|
+
from defuser.utils.common import env_flag, warn_if_public_api_transformers_unsupported
|
|
7
|
+
from logbar import LogBar
|
|
7
8
|
|
|
8
9
|
DEBUG_ON = env_flag("DEBUG")
|
|
9
10
|
|
|
10
|
-
|
|
11
|
+
logger = LogBar(__name__)
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
def convert_model(*args, **kwargs) -> bool:
|
|
11
15
|
"""Lazily import conversion entrypoint to avoid import-time cycles."""
|
|
16
|
+
if warn_if_public_api_transformers_unsupported("convert_model()", logger):
|
|
17
|
+
return False
|
|
18
|
+
|
|
12
19
|
from .defuser import convert_model as _convert_model
|
|
13
20
|
|
|
14
21
|
return _convert_model(*args, **kwargs)
|
|
15
22
|
|
|
16
23
|
|
|
17
|
-
def replace_fused_blocks(*args, **kwargs):
|
|
24
|
+
def replace_fused_blocks(*args, **kwargs) -> bool:
|
|
18
25
|
"""Lazily import conversion entrypoint to avoid import-time cycles."""
|
|
26
|
+
if warn_if_public_api_transformers_unsupported("replace_fused_blocks()", logger):
|
|
27
|
+
return False
|
|
28
|
+
|
|
19
29
|
from .defuser import replace_fused_blocks as _replace_fused_blocks
|
|
20
30
|
|
|
21
31
|
return _replace_fused_blocks(*args, **kwargs)
|
|
@@ -1,3 +1,7 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: 2026 ModelCloud.ai
|
|
2
|
+
# SPDX-FileCopyrightText: 2026 qubitium@modelcloud.ai
|
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
# Contact: qubitium@modelcloud.ai, x.com/qubitium
|
|
1
5
|
import torch
|
|
2
6
|
from transformers.core_model_loading import Chunk, Concatenate, ConversionOps, MergeModulelist
|
|
3
7
|
|
|
@@ -36,9 +36,26 @@ def get_checkpoint_conversion_mapping(model_type):
|
|
|
36
36
|
|
|
37
37
|
|
|
38
38
|
class PatchError(Exception):
|
|
39
|
+
"""Raised when Defuser cannot patch a registered Transformers class."""
|
|
40
|
+
|
|
39
41
|
pass
|
|
40
42
|
|
|
41
43
|
|
|
44
|
+
def _has_prebuilt_replacements(model: nn.Module, model_type: str) -> bool:
|
|
45
|
+
"""Detect models that were already instantiated with registry-backed replacements."""
|
|
46
|
+
replacement_paths = MODEL_CONFIG[model_type].get(PATCH.REPLACE_MODULE, [])
|
|
47
|
+
replacement_class_paths = {custom_path for _, custom_path in replacement_paths}
|
|
48
|
+
if not replacement_class_paths:
|
|
49
|
+
return False
|
|
50
|
+
|
|
51
|
+
for module in model.modules():
|
|
52
|
+
class_path = f"{module.__class__.__module__}.{module.__class__.__name__}"
|
|
53
|
+
if class_path in replacement_class_paths:
|
|
54
|
+
return True
|
|
55
|
+
|
|
56
|
+
return False
|
|
57
|
+
|
|
58
|
+
|
|
42
59
|
def replace_fused_blocks(model_type: str) -> bool:
|
|
43
60
|
"""Patch supported HF model classes so future loads instantiate defused blocks."""
|
|
44
61
|
if warn_if_public_api_transformers_unsupported("replace_fused_blocks()", logger):
|
|
@@ -202,9 +219,10 @@ def convert_model(
|
|
|
202
219
|
|
|
203
220
|
apply_model_patches(model, max_layers=max_layers, filter_rules=filter)
|
|
204
221
|
|
|
205
|
-
#
|
|
206
|
-
#
|
|
207
|
-
|
|
222
|
+
# Full models patched at construction time already contain the defused
|
|
223
|
+
# replacement modules, but standalone experts from those model types can
|
|
224
|
+
# still use runtime tensor defusion.
|
|
225
|
+
if _has_prebuilt_replacements(model, model.config.model_type):
|
|
208
226
|
return False
|
|
209
227
|
|
|
210
228
|
# Perform runtime defusing of fused projections
|
|
@@ -135,6 +135,12 @@ MODEL_CONFIG = {
|
|
|
135
135
|
},
|
|
136
136
|
"glm4_moe_lite": {
|
|
137
137
|
"min_transformers_version": MIN_SUPPORTED_TRANSFORMERS_VERSION,
|
|
138
|
+
PATCH.REPLACE_MODULE: [
|
|
139
|
+
(
|
|
140
|
+
"transformers.models.glm4_moe_lite.modeling_glm4_moe_lite.Glm4MoeLiteMoE",
|
|
141
|
+
"defuser.modeling.unfused_moe.glm4_moe_lite.LinearGlm4MoeLiteMoE",
|
|
142
|
+
)
|
|
143
|
+
],
|
|
138
144
|
},
|
|
139
145
|
"glm4v": {
|
|
140
146
|
"min_transformers_version": MIN_SUPPORTED_TRANSFORMERS_VERSION,
|
|
@@ -292,28 +292,36 @@ def _install_instance_forward(module: nn.Module, implementation: str) -> None:
|
|
|
292
292
|
def _detect_expert_projections(module: nn.Module) -> dict[str, dict]:
|
|
293
293
|
"""Detect which expert projections exist in the module.
|
|
294
294
|
|
|
295
|
-
This function scans the module for any 3D
|
|
296
|
-
It first checks known projection names, then discovers any
|
|
295
|
+
This function scans the module for any registered 3D Parameter / Tensor
|
|
296
|
+
attributes. It first checks known projection names, then discovers any
|
|
297
|
+
unknown 3D registered tensors.
|
|
297
298
|
|
|
298
299
|
Returns:
|
|
299
300
|
Dict mapping projection names to their config, only for projections that exist
|
|
300
|
-
as 3D
|
|
301
|
+
as 3D registered Parameter / Tensor in the module.
|
|
301
302
|
"""
|
|
302
303
|
detected = {}
|
|
304
|
+
# Only inspect registered tensors here. Scanning arbitrary attributes can
|
|
305
|
+
# trigger unrelated properties such as Transformers' `loss_function`.
|
|
306
|
+
local_tensors = {
|
|
307
|
+
name: tensor
|
|
308
|
+
for registry in (module._parameters, module._buffers)
|
|
309
|
+
for name, tensor in registry.items()
|
|
310
|
+
if isinstance(tensor, torch.Tensor)
|
|
311
|
+
}
|
|
303
312
|
|
|
304
313
|
# First, check known projection patterns
|
|
305
314
|
for proj_name, config in KNOWN_PROJECTION_PATTERNS.items():
|
|
306
|
-
param =
|
|
307
|
-
if param is not None and
|
|
315
|
+
param = local_tensors.get(proj_name)
|
|
316
|
+
if param is not None and param.dim() == 3:
|
|
308
317
|
detected[proj_name] = config
|
|
309
318
|
|
|
310
|
-
# If no known patterns found, scan for any 3D
|
|
319
|
+
# If no known patterns found, scan for any 3D registered tensor (future-proofing)
|
|
311
320
|
if not detected:
|
|
312
|
-
for attr_name in
|
|
321
|
+
for attr_name, param in local_tensors.items():
|
|
313
322
|
if attr_name.startswith("_"):
|
|
314
323
|
continue
|
|
315
|
-
param
|
|
316
|
-
if param is not None and isinstance(param, nn.Parameter) and param.dim() == 3:
|
|
324
|
+
if param is not None and param.dim() == 3:
|
|
317
325
|
# Use default config for unknown projections
|
|
318
326
|
if DEBUG_ON: logger.debug(f"Discovered unknown 3D projection: {attr_name}")
|
|
319
327
|
detected[attr_name] = {"is_input_proj": True, "output_multiplier": 1}
|
|
@@ -321,6 +329,29 @@ def _detect_expert_projections(module: nn.Module) -> dict[str, dict]:
|
|
|
321
329
|
return detected
|
|
322
330
|
|
|
323
331
|
|
|
332
|
+
def _get_registered_tensor(module: nn.Module, name: str) -> torch.Tensor | None:
|
|
333
|
+
"""Return a direct registered parameter or buffer without touching properties."""
|
|
334
|
+
|
|
335
|
+
tensor = module._parameters.get(name)
|
|
336
|
+
if isinstance(tensor, torch.Tensor):
|
|
337
|
+
return tensor
|
|
338
|
+
|
|
339
|
+
tensor = module._buffers.get(name)
|
|
340
|
+
if isinstance(tensor, torch.Tensor):
|
|
341
|
+
return tensor
|
|
342
|
+
|
|
343
|
+
return None
|
|
344
|
+
|
|
345
|
+
|
|
346
|
+
def _set_registered_tensor_like(module: nn.Module, name: str, tensor: torch.Tensor, source: torch.Tensor) -> None:
|
|
347
|
+
"""Register ``tensor`` using the same parameter-vs-buffer kind as ``source``."""
|
|
348
|
+
|
|
349
|
+
if isinstance(source, nn.Parameter):
|
|
350
|
+
module.register_parameter(name, nn.Parameter(tensor, requires_grad=source.requires_grad))
|
|
351
|
+
else:
|
|
352
|
+
module.register_buffer(name, tensor)
|
|
353
|
+
|
|
354
|
+
|
|
324
355
|
def _experts_supports_decorator(module: nn.Module) -> bool:
|
|
325
356
|
"""Check if experts module supports @use_experts_implementation decorator.
|
|
326
357
|
|
|
@@ -334,11 +365,11 @@ def _experts_supports_decorator(module: nn.Module) -> bool:
|
|
|
334
365
|
return hasattr(forward_method, "__wrapped__")
|
|
335
366
|
|
|
336
367
|
|
|
337
|
-
def _infer_dimensions(param:
|
|
368
|
+
def _infer_dimensions(param: torch.Tensor, config: dict, is_transposed: bool) -> tuple[int, int]:
|
|
338
369
|
"""Infer input and output dimensions for a projection.
|
|
339
370
|
|
|
340
371
|
Args:
|
|
341
|
-
param: The 3D
|
|
372
|
+
param: The 3D projection tensor (num_experts, dim1, dim2)
|
|
342
373
|
config: Projection config with is_input_proj and output_multiplier
|
|
343
374
|
is_transposed: Whether weights are stored transposed
|
|
344
375
|
|
|
@@ -370,7 +401,7 @@ def _unfuse_single_projection(
|
|
|
370
401
|
dtype: torch.dtype,
|
|
371
402
|
target_device: torch.device,
|
|
372
403
|
) -> list | None:
|
|
373
|
-
"""Unfuse a single projection from 3D
|
|
404
|
+
"""Unfuse a single projection from a 3D registered tensor to Linear layers.
|
|
374
405
|
|
|
375
406
|
Optimized to keep peak device memory low while preserving the module's
|
|
376
407
|
original device placement:
|
|
@@ -390,8 +421,8 @@ def _unfuse_single_projection(
|
|
|
390
421
|
Returns:
|
|
391
422
|
List of Linear layers, or None if projection doesn't exist
|
|
392
423
|
"""
|
|
393
|
-
param =
|
|
394
|
-
if param is None or
|
|
424
|
+
param = _get_registered_tensor(module, proj_name)
|
|
425
|
+
if param is None or param.dim() != 3:
|
|
395
426
|
return None
|
|
396
427
|
|
|
397
428
|
# Get projection config
|
|
@@ -402,17 +433,17 @@ def _unfuse_single_projection(
|
|
|
402
433
|
|
|
403
434
|
# Check for bias
|
|
404
435
|
bias_name = f"{proj_name}_bias"
|
|
405
|
-
bias_param =
|
|
406
|
-
has_bias = bias_param
|
|
436
|
+
bias_param = _get_registered_tensor(module, bias_name)
|
|
437
|
+
has_bias = isinstance(bias_param, torch.Tensor)
|
|
407
438
|
|
|
408
439
|
source_device = param.device
|
|
409
440
|
is_meta = source_device.type == "meta"
|
|
410
|
-
weight_requires_grad = param.requires_grad
|
|
441
|
+
weight_requires_grad = param.requires_grad if isinstance(param, nn.Parameter) else False
|
|
411
442
|
|
|
412
443
|
# Prepare weight slices on CPU in batch (single D2H transfer + batch transpose)
|
|
413
444
|
if not is_meta:
|
|
414
445
|
# Single transfer: GPU -> CPU (or no-op if already on CPU)
|
|
415
|
-
weights_cpu = param.
|
|
446
|
+
weights_cpu = param.detach().cpu() # (num_experts, dim1, dim2)
|
|
416
447
|
if is_transposed:
|
|
417
448
|
# Batch transpose: (num_experts, in, out) -> (num_experts, out, in)
|
|
418
449
|
weights_cpu = weights_cpu.transpose(1, 2)
|
|
@@ -425,12 +456,12 @@ def _unfuse_single_projection(
|
|
|
425
456
|
weight_slices = weights_cpu.unbind(0)
|
|
426
457
|
|
|
427
458
|
if has_bias:
|
|
428
|
-
bias_cpu = bias_param.
|
|
459
|
+
bias_cpu = bias_param.detach().cpu()
|
|
429
460
|
# Ensure contiguous — bias may come from a chunk split (Phase 1)
|
|
430
461
|
if not bias_cpu.is_contiguous():
|
|
431
462
|
bias_cpu = bias_cpu.contiguous()
|
|
432
463
|
bias_slices = bias_cpu.unbind(0)
|
|
433
|
-
bias_requires_grad = bias_param.requires_grad
|
|
464
|
+
bias_requires_grad = bias_param.requires_grad if isinstance(bias_param, nn.Parameter) else False
|
|
434
465
|
|
|
435
466
|
# Drop the original fused parameter before allocating the defused
|
|
436
467
|
# per-expert linears back on the original device.
|
|
@@ -588,7 +619,7 @@ def _unfuse_experts_weights_inplace(
|
|
|
588
619
|
|
|
589
620
|
# Get first projection to determine num_experts and layout
|
|
590
621
|
first_proj_name = next(iter(detected_projections))
|
|
591
|
-
first_param =
|
|
622
|
+
first_param = _get_registered_tensor(module, first_proj_name)
|
|
592
623
|
num_experts = first_param.shape[0]
|
|
593
624
|
|
|
594
625
|
# Detect if transposed
|
|
@@ -608,15 +639,15 @@ def _unfuse_experts_weights_inplace(
|
|
|
608
639
|
dtype = first_param.dtype
|
|
609
640
|
target_device = first_param.device if first_param.device.type != "meta" else "cpu"
|
|
610
641
|
|
|
611
|
-
# Phase 1: Split fused projections (e.g., gate_up_proj -> gate_proj + up_proj) into separate 3D
|
|
642
|
+
# Phase 1: Split fused projections (e.g., gate_up_proj -> gate_proj + up_proj) into separate 3D tensors
|
|
612
643
|
extra_projections = {}
|
|
613
644
|
fused_to_remove = []
|
|
614
645
|
for proj_name, config in detected_projections.items():
|
|
615
646
|
split_into = config.get("split_into")
|
|
616
647
|
if not split_into:
|
|
617
648
|
continue
|
|
618
|
-
param =
|
|
619
|
-
if param is None or
|
|
649
|
+
param = _get_registered_tensor(module, proj_name)
|
|
650
|
+
if param is None or param.dim() != 3:
|
|
620
651
|
continue
|
|
621
652
|
# Split along output dimension
|
|
622
653
|
split_dim = 2 if is_transposed else 1
|
|
@@ -624,17 +655,17 @@ def _unfuse_experts_weights_inplace(
|
|
|
624
655
|
|
|
625
656
|
# Also split bias if present (e.g., gate_up_proj_bias -> gate_proj_bias + up_proj_bias)
|
|
626
657
|
bias_name = f"{proj_name}_bias"
|
|
627
|
-
bias_param =
|
|
658
|
+
bias_param = _get_registered_tensor(module, bias_name)
|
|
628
659
|
bias_splits = None
|
|
629
|
-
if
|
|
660
|
+
if isinstance(bias_param, torch.Tensor) and bias_param.dim() == 2:
|
|
630
661
|
bias_splits = bias_param.chunk(len(split_into), dim=1)
|
|
631
662
|
|
|
632
663
|
for i, (split_name, split_param) in enumerate(zip(split_into, split_params)):
|
|
633
664
|
# Avoid .contiguous() here — _unfuse_single_projection will handle it
|
|
634
665
|
# during batch transpose/unbind, saving a full-tensor copy
|
|
635
|
-
|
|
666
|
+
_set_registered_tensor_like(module, split_name, split_param, param)
|
|
636
667
|
if bias_splits is not None:
|
|
637
|
-
|
|
668
|
+
_set_registered_tensor_like(module, f"{split_name}_bias", bias_splits[i], bias_param)
|
|
638
669
|
extra_projections[split_name] = KNOWN_PROJECTION_PATTERNS.get(
|
|
639
670
|
split_name, {"is_input_proj": True, "output_multiplier": 1}
|
|
640
671
|
)
|
|
@@ -649,7 +680,7 @@ def _unfuse_experts_weights_inplace(
|
|
|
649
680
|
del detected_projections[name]
|
|
650
681
|
detected_projections.update(extra_projections)
|
|
651
682
|
|
|
652
|
-
# Phase 2: Unfuse all 3D
|
|
683
|
+
# Phase 2: Unfuse all 3D tensors into per-expert Linear layers
|
|
653
684
|
proj_linears = {} # proj_name -> [Linear_expert0, Linear_expert1, ...]
|
|
654
685
|
for proj_name in detected_projections:
|
|
655
686
|
linears = _unfuse_single_projection(module, proj_name, num_experts, is_transposed, dtype, target_device)
|
|
@@ -0,0 +1,100 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: 2026 ModelCloud.ai
|
|
2
|
+
# SPDX-FileCopyrightText: 2026 qubitium@modelcloud.ai
|
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
# Contact: qubitium@modelcloud.ai, x.com/qubitium
|
|
5
|
+
|
|
6
|
+
# Adapted from intel/auto-round
|
|
7
|
+
# at https://github.com/intel/auto-round/blob/main/auto_round/modeling/unfused_moe/glm_moe_light.py
|
|
8
|
+
|
|
9
|
+
import torch
|
|
10
|
+
import torch.nn as nn
|
|
11
|
+
|
|
12
|
+
class LinearGlm4MoeLiteMoE(nn.Module):
|
|
13
|
+
"""
|
|
14
|
+
A mixed expert module containing shared experts.
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
def __init__(self, config):
|
|
18
|
+
super().__init__()
|
|
19
|
+
self.config = config
|
|
20
|
+
self.num_experts = config.num_local_experts
|
|
21
|
+
from transformers.models.glm4_moe_lite.modeling_glm4_moe_lite import Glm4MoeLiteMLP, Glm4MoeLiteTopkRouter
|
|
22
|
+
|
|
23
|
+
self.experts = nn.ModuleList(
|
|
24
|
+
[Glm4MoeLiteMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(self.num_experts)]
|
|
25
|
+
)
|
|
26
|
+
|
|
27
|
+
self.gate = Glm4MoeLiteTopkRouter(config)
|
|
28
|
+
self.shared_experts = Glm4MoeLiteMLP(
|
|
29
|
+
config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts
|
|
30
|
+
)
|
|
31
|
+
self.n_routed_experts = config.n_routed_experts
|
|
32
|
+
self.n_group = config.n_group
|
|
33
|
+
self.topk_group = config.topk_group
|
|
34
|
+
self.norm_topk_prob = config.norm_topk_prob
|
|
35
|
+
self.routed_scaling_factor = config.routed_scaling_factor
|
|
36
|
+
self.top_k = config.num_experts_per_tok
|
|
37
|
+
|
|
38
|
+
def experts_forward(
|
|
39
|
+
self,
|
|
40
|
+
hidden_states: torch.Tensor,
|
|
41
|
+
top_k_index: torch.Tensor,
|
|
42
|
+
top_k_weights: torch.Tensor,
|
|
43
|
+
) -> torch.Tensor:
|
|
44
|
+
""" """
|
|
45
|
+
final_hidden_states = torch.zeros_like(hidden_states)
|
|
46
|
+
with torch.no_grad():
|
|
47
|
+
expert_mask = torch.nn.functional.one_hot(top_k_index, num_classes=self.num_experts)
|
|
48
|
+
expert_mask = expert_mask.permute(2, 1, 0)
|
|
49
|
+
expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
|
|
50
|
+
|
|
51
|
+
for expert_idx in expert_hit:
|
|
52
|
+
expert_idx = expert_idx[0]
|
|
53
|
+
if expert_idx == self.num_experts:
|
|
54
|
+
continue
|
|
55
|
+
top_k_pos, token_idx = torch.where(expert_mask[expert_idx])
|
|
56
|
+
current_state = hidden_states[token_idx]
|
|
57
|
+
expert_layer = self.experts[expert_idx]
|
|
58
|
+
current_hidden_states = expert_layer(current_state)
|
|
59
|
+
# gate, up = nn.functional.linear(current_state, self.gate_up_proj[expert_idx]).chunk(2, dim=-1)
|
|
60
|
+
# current_hidden_states = self.act_fn(gate) * up
|
|
61
|
+
# current_hidden_states = nn.functional.linear(current_hidden_states, self.down_proj[expert_idx])
|
|
62
|
+
current_hidden_states = current_hidden_states * top_k_weights[token_idx, top_k_pos, None]
|
|
63
|
+
final_hidden_states.index_add_(0, token_idx, current_hidden_states.to(final_hidden_states.dtype))
|
|
64
|
+
|
|
65
|
+
return final_hidden_states
|
|
66
|
+
|
|
67
|
+
def route_tokens_to_experts(self, router_logits):
|
|
68
|
+
router_logits = router_logits.sigmoid()
|
|
69
|
+
router_logits_for_choice = router_logits + self.gate.e_score_correction_bias
|
|
70
|
+
group_scores = (
|
|
71
|
+
router_logits_for_choice.view(-1, self.n_group, self.n_routed_experts // self.n_group)
|
|
72
|
+
.topk(2, dim=-1)[0]
|
|
73
|
+
.sum(dim=-1)
|
|
74
|
+
)
|
|
75
|
+
group_idx = torch.topk(group_scores, k=self.topk_group, dim=-1, sorted=False)[1]
|
|
76
|
+
group_mask = torch.zeros_like(group_scores)
|
|
77
|
+
group_mask.scatter_(1, group_idx, 1)
|
|
78
|
+
score_mask = (
|
|
79
|
+
group_mask.unsqueeze(-1)
|
|
80
|
+
.expand(-1, self.n_group, self.n_routed_experts // self.n_group)
|
|
81
|
+
.reshape(-1, self.n_routed_experts)
|
|
82
|
+
)
|
|
83
|
+
scores_for_choice = router_logits_for_choice.masked_fill(~score_mask.bool(), 0.0)
|
|
84
|
+
topk_indices = torch.topk(scores_for_choice, k=self.top_k, dim=-1, sorted=False)[1]
|
|
85
|
+
topk_weights = router_logits.gather(1, topk_indices)
|
|
86
|
+
if self.norm_topk_prob:
|
|
87
|
+
denominator = topk_weights.sum(dim=-1, keepdim=True) + 1e-20
|
|
88
|
+
topk_weights /= denominator
|
|
89
|
+
topk_weights = topk_weights * self.routed_scaling_factor
|
|
90
|
+
return topk_indices, topk_weights
|
|
91
|
+
|
|
92
|
+
def forward(self, hidden_states):
|
|
93
|
+
residuals = hidden_states
|
|
94
|
+
orig_shape = hidden_states.shape
|
|
95
|
+
router_logits = self.gate(hidden_states)
|
|
96
|
+
topk_indices, topk_weights = self.route_tokens_to_experts(router_logits)
|
|
97
|
+
hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
|
|
98
|
+
hidden_states = self.experts_forward(hidden_states, topk_indices, topk_weights).view(*orig_shape)
|
|
99
|
+
hidden_states = hidden_states + self.shared_experts(residuals)
|
|
100
|
+
return hidden_states
|
|
@@ -9,7 +9,7 @@ build-backend = "setuptools.build_meta"
|
|
|
9
9
|
|
|
10
10
|
[project]
|
|
11
11
|
name = "Defuser"
|
|
12
|
-
version = "0.0.
|
|
12
|
+
version = "0.0.19"
|
|
13
13
|
description = "Model defuser helper for HF Transformers."
|
|
14
14
|
readme = "README.md"
|
|
15
15
|
requires-python = ">=3.9"
|
|
@@ -25,6 +25,7 @@ classifiers = [
|
|
|
25
25
|
"Programming Language :: Python :: 3.12",
|
|
26
26
|
"Programming Language :: Python :: 3.13",
|
|
27
27
|
"Programming Language :: Python :: 3.14",
|
|
28
|
+
"Programming Language :: Python :: 3.15",
|
|
28
29
|
"Programming Language :: C",
|
|
29
30
|
"Topic :: Software Development :: Libraries",
|
|
30
31
|
"Topic :: Text Processing :: Linguistic"
|
defuser-0.0.19/setup.py
ADDED
|
@@ -17,12 +17,16 @@ from defuser.model_registry import MODEL_CONFIG
|
|
|
17
17
|
|
|
18
18
|
|
|
19
19
|
class _DummyLayer(nn.Module):
|
|
20
|
+
"""Wrap one candidate module so convert_model sees a layer-shaped path."""
|
|
21
|
+
|
|
20
22
|
def __init__(self, module: nn.Module):
|
|
21
23
|
super().__init__()
|
|
22
24
|
self.module = module
|
|
23
25
|
|
|
24
26
|
|
|
25
27
|
class _DummyModel(nn.Module):
|
|
28
|
+
"""Minimal model shell for exercising runtime conversion helpers."""
|
|
29
|
+
|
|
26
30
|
def __init__(self, model_type: str, module: nn.Module):
|
|
27
31
|
super().__init__()
|
|
28
32
|
self.config = getattr(module, "config", None)
|
|
@@ -33,10 +37,21 @@ class _DummyModel(nn.Module):
|
|
|
33
37
|
|
|
34
38
|
|
|
35
39
|
def _load(module_path: str, attr_name: str):
|
|
40
|
+
"""Import one symbol lazily so candidate tables stay compact."""
|
|
36
41
|
return getattr(import_module(module_path), attr_name)
|
|
37
42
|
|
|
38
43
|
|
|
44
|
+
def _normalize_attr_value(obj, attr: str, value):
|
|
45
|
+
"""Preserve list-valued config shapes when tests shrink scalar attributes."""
|
|
46
|
+
current = getattr(obj, attr)
|
|
47
|
+
if isinstance(current, list) and not isinstance(value, list):
|
|
48
|
+
width = len(current) or 1
|
|
49
|
+
return [value for _ in range(width)]
|
|
50
|
+
return value
|
|
51
|
+
|
|
52
|
+
|
|
39
53
|
def _build_config(case: dict):
|
|
54
|
+
"""Build a tiny but valid config for one candidate module case."""
|
|
40
55
|
config = _load(case["config_module"], case["config_name"])()
|
|
41
56
|
sub_attr = case.get("sub_attr")
|
|
42
57
|
if sub_attr is not None:
|
|
@@ -73,24 +88,29 @@ def _build_config(case: dict):
|
|
|
73
88
|
"dropout_rate": 0.0,
|
|
74
89
|
}.items():
|
|
75
90
|
if hasattr(config, attr):
|
|
76
|
-
setattr(config, attr, value)
|
|
91
|
+
setattr(config, attr, _normalize_attr_value(config, attr, value))
|
|
77
92
|
for attr, value in case.get("config_updates", {}).items():
|
|
78
93
|
if hasattr(config, attr):
|
|
79
|
-
setattr(config, attr, value)
|
|
94
|
+
setattr(config, attr, _normalize_attr_value(config, attr, value))
|
|
80
95
|
return config
|
|
81
96
|
|
|
82
97
|
|
|
83
98
|
def _build_module(case: dict) -> nn.Module:
|
|
99
|
+
"""Instantiate one candidate module in the format its constructor expects."""
|
|
84
100
|
module_cls = _load(case["module_path"], case["class_name"])
|
|
85
101
|
kind = case.get("kind", "config")
|
|
86
102
|
if kind == "parallel":
|
|
87
103
|
return module_cls(case["num_experts"], case["input_size"], case["output_size"]).eval()
|
|
88
104
|
if kind == "zamba2":
|
|
89
105
|
return module_cls(_build_config(case), num_fwd_mem_blocks=1, block_id=0).eval()
|
|
90
|
-
|
|
106
|
+
config = _build_config(case)
|
|
107
|
+
if case["model_type"] == "ernie4_5_vl_moe" and isinstance(getattr(config, "moe_intermediate_size", None), list):
|
|
108
|
+
return module_cls(config, intermediate_size=config.moe_intermediate_size[0]).eval()
|
|
109
|
+
return module_cls(config).eval()
|
|
91
110
|
|
|
92
111
|
|
|
93
112
|
def _wrapped_model(case: dict) -> tuple[_DummyModel, nn.Module]:
|
|
113
|
+
"""Seed one module deterministically and wrap it as a dummy model."""
|
|
94
114
|
module = _build_module(case)
|
|
95
115
|
generator = torch.Generator(device="cpu").manual_seed(0)
|
|
96
116
|
with torch.no_grad():
|
|
@@ -101,10 +121,12 @@ def _wrapped_model(case: dict) -> tuple[_DummyModel, nn.Module]:
|
|
|
101
121
|
|
|
102
122
|
|
|
103
123
|
def _patched_module(model: _DummyModel) -> nn.Module:
|
|
124
|
+
"""Return the converted test module from the dummy model wrapper."""
|
|
104
125
|
return model.layers[0].module
|
|
105
126
|
|
|
106
127
|
|
|
107
128
|
def _assert_expert_container(module: nn.Module, attrs: tuple[str, ...]) -> None:
|
|
129
|
+
"""Verify that expert ``0`` exposes the expected defused projections."""
|
|
108
130
|
assert hasattr(module, "0")
|
|
109
131
|
expert0 = getattr(module, "0")
|
|
110
132
|
for attr in attrs:
|
|
@@ -112,6 +134,7 @@ def _assert_expert_container(module: nn.Module, attrs: tuple[str, ...]) -> None:
|
|
|
112
134
|
|
|
113
135
|
|
|
114
136
|
def _standard_hidden(case: dict) -> torch.Tensor:
|
|
137
|
+
"""Generate default hidden states for sparse-MoE candidate coverage tests."""
|
|
115
138
|
return torch.randn(case.get("hidden_shape", (5, case["input_dim"])), dtype=torch.float32)
|
|
116
139
|
|
|
117
140
|
|
|
@@ -553,11 +576,13 @@ ALL_CANDIDATE_MODEL_TYPES = {
|
|
|
553
576
|
|
|
554
577
|
|
|
555
578
|
def test_model_registry_covers_all_scanned_candidates():
|
|
579
|
+
"""Every scanned candidate model type should be represented in the registry."""
|
|
556
580
|
assert ALL_CANDIDATE_MODEL_TYPES.issubset(MODEL_CONFIG)
|
|
557
581
|
|
|
558
582
|
|
|
559
583
|
@pytest.mark.parametrize("case", STANDARD_MOE_CASES, ids=[case["model_type"] for case in STANDARD_MOE_CASES])
|
|
560
584
|
def test_standard_moe_candidates_convert_and_preserve_forward(case):
|
|
585
|
+
"""Standard sparse-MoE candidates should convert without changing outputs."""
|
|
561
586
|
torch.manual_seed(0)
|
|
562
587
|
model, original_module = _wrapped_model(case)
|
|
563
588
|
hidden_states = _standard_hidden(case)
|
|
@@ -593,6 +618,7 @@ def test_standard_moe_candidates_convert_and_preserve_forward(case):
|
|
|
593
618
|
|
|
594
619
|
@pytest.mark.parametrize("case", PARALLEL_CASES, ids=[case["model_type"] for case in PARALLEL_CASES])
|
|
595
620
|
def test_parallel_expert_candidates_convert_and_preserve_forward(case):
|
|
621
|
+
"""Parallel expert containers should convert without changing outputs."""
|
|
596
622
|
torch.manual_seed(0)
|
|
597
623
|
model, original_module = _wrapped_model(case)
|
|
598
624
|
expert_size = [2, 1, 0, 3]
|
|
@@ -616,6 +642,7 @@ def test_parallel_expert_candidates_convert_and_preserve_forward(case):
|
|
|
616
642
|
|
|
617
643
|
@pytest.mark.parametrize("case", DENSE_CASES, ids=[case["label"] for case in DENSE_CASES])
|
|
618
644
|
def test_dense_candidates_convert_and_preserve_forward(case):
|
|
645
|
+
"""Dense split-gate MLP candidates should preserve forward results after conversion."""
|
|
619
646
|
torch.manual_seed(0)
|
|
620
647
|
model, original_module = _wrapped_model(case)
|
|
621
648
|
hidden_states = torch.randn(3, case["hidden_size"], dtype=torch.float32)
|
|
@@ -644,6 +671,7 @@ def test_dense_candidates_convert_and_preserve_forward(case):
|
|
|
644
671
|
|
|
645
672
|
|
|
646
673
|
def test_runtime_model_patches_respect_max_layers():
|
|
674
|
+
"""Runtime-only patches should stop once ``max_layers`` is reached."""
|
|
647
675
|
module0 = _build_module(next(case for case in DENSE_CASES if case["label"] == "phi3"))
|
|
648
676
|
module1 = _build_module(next(case for case in DENSE_CASES if case["label"] == "phi3"))
|
|
649
677
|
|
|
@@ -10,6 +10,11 @@ from safetensors.torch import save_file
|
|
|
10
10
|
from torch import nn
|
|
11
11
|
from transformers.core_model_loading import WeightConverter
|
|
12
12
|
from transformers.models.glm4_moe.modeling_glm4_moe import Glm4MoeConfig, Glm4MoeForCausalLM, Glm4MoeMoE
|
|
13
|
+
from transformers.models.glm4_moe_lite.configuration_glm4_moe_lite import Glm4MoeLiteConfig
|
|
14
|
+
from transformers.models.glm4_moe_lite.modeling_glm4_moe_lite import (
|
|
15
|
+
Glm4MoeLiteForCausalLM,
|
|
16
|
+
Glm4MoeLiteMoE,
|
|
17
|
+
)
|
|
13
18
|
from transformers.models.glm4v.configuration_glm4v import Glm4vConfig
|
|
14
19
|
from transformers.models.glm4v.modeling_glm4v import Glm4vForConditionalGeneration
|
|
15
20
|
from transformers.models.mixtral.configuration_mixtral import MixtralConfig
|
|
@@ -48,6 +53,7 @@ from defuser.checkpoint_ops import OwnedChunk, SplitFusedExpertDownProj, SplitFu
|
|
|
48
53
|
from defuser.model_registry import MODEL_CONFIG, PATCH
|
|
49
54
|
from defuser.modeling.replace_modules import ReplacementModuleBase, apply_replacements, materialize_model
|
|
50
55
|
from defuser.modeling.unfused_moe.glm4_moe import LinearGlm4MoeMoE
|
|
56
|
+
from defuser.modeling.unfused_moe.glm4_moe_lite import LinearGlm4MoeLiteMoE
|
|
51
57
|
from defuser.modeling.unfused_moe.mixtral import LinearMixtralSparseMoeBlock
|
|
52
58
|
from defuser.modeling.unfused_moe.qwen2_moe import LinearQwen2MoeSparseMoeBlock
|
|
53
59
|
from defuser.modeling.unfused_moe.qwen3_moe import LinearQwen3MoeSparseMoeBlock
|
|
@@ -61,9 +67,9 @@ from defuser.utils.common import MIN_SUPPORTED_TRANSFORMERS_VERSION
|
|
|
61
67
|
|
|
62
68
|
|
|
63
69
|
|
|
64
|
-
def _tiny_moe_config(config_cls):
|
|
70
|
+
def _tiny_moe_config(config_cls, num_hidden_layers: int=1):
|
|
65
71
|
return config_cls(
|
|
66
|
-
num_hidden_layers=
|
|
72
|
+
num_hidden_layers=num_hidden_layers,
|
|
67
73
|
hidden_size=64,
|
|
68
74
|
intermediate_size=128,
|
|
69
75
|
moe_intermediate_size=32,
|
|
@@ -78,6 +84,7 @@ def _tiny_moe_config(config_cls):
|
|
|
78
84
|
|
|
79
85
|
def _tiny_qwen3_omni_config():
|
|
80
86
|
return Qwen3OmniMoeConfig(
|
|
87
|
+
initializer_range=0.02,
|
|
81
88
|
enable_audio_output=False,
|
|
82
89
|
thinker_config={
|
|
83
90
|
"text_config": {
|
|
@@ -936,6 +943,29 @@ def test_defused_models_preserve_output_router_logits_capture():
|
|
|
936
943
|
assert len(outputs.router_logits) == 1
|
|
937
944
|
assert outputs.router_logits[0].shape == (3, model.config.num_experts)
|
|
938
945
|
|
|
946
|
+
def test_glm4_moe_lite():
|
|
947
|
+
model_type = "glm4_moe_lite"
|
|
948
|
+
replace_fused_blocks(model_type)
|
|
949
|
+
|
|
950
|
+
model = Glm4MoeLiteForCausalLM(_tiny_moe_config(Glm4MoeLiteConfig, num_hidden_layers=2))
|
|
951
|
+
assert model.config.model_type == model_type
|
|
952
|
+
|
|
953
|
+
# In GLM4-MoE-Lite, the `mlp.experts` module is present only starting from the second layer.
|
|
954
|
+
converted = convert_model(model, max_layers=2)
|
|
955
|
+
assert not converted
|
|
956
|
+
|
|
957
|
+
_assert_unfused_expert_module(model.model.layers[1].mlp.experts)
|
|
958
|
+
|
|
959
|
+
|
|
960
|
+
def test_glm4_moe_lite_defused_forward_matches_fused_math():
|
|
961
|
+
config = _tiny_moe_config(Glm4MoeLiteConfig)
|
|
962
|
+
hidden_states = torch.randn(2, 3, config.hidden_size, dtype=torch.float32)
|
|
963
|
+
|
|
964
|
+
_assert_sparse_moe_defused_matches_fused_math(
|
|
965
|
+
Glm4MoeLiteMoE(config),
|
|
966
|
+
LinearGlm4MoeLiteMoE(config),
|
|
967
|
+
hidden_states,
|
|
968
|
+
)
|
|
939
969
|
|
|
940
970
|
def test_glm4v_checkpoint_mapping_splits_gate_up_proj():
|
|
941
971
|
from defuser.defuser import get_checkpoint_conversion_mapping
|
|
@@ -15,16 +15,28 @@ from defuser.model_registry import MODEL_CONFIG
|
|
|
15
15
|
|
|
16
16
|
|
|
17
17
|
def _load(module_path: str, attr_name: str):
|
|
18
|
+
"""Import one symbol lazily so the model case table stays readable."""
|
|
18
19
|
return getattr(import_module(module_path), attr_name)
|
|
19
20
|
|
|
20
21
|
|
|
22
|
+
def _normalize_attr_value(obj, attr: str, value):
|
|
23
|
+
"""Preserve list-valued config fields when shrinking configs for meta tests."""
|
|
24
|
+
current = getattr(obj, attr)
|
|
25
|
+
if isinstance(current, list) and not isinstance(value, list):
|
|
26
|
+
width = len(current) or 1
|
|
27
|
+
return [value for _ in range(width)]
|
|
28
|
+
return value
|
|
29
|
+
|
|
30
|
+
|
|
21
31
|
def _set_if_has(obj, **kwargs) -> None:
|
|
32
|
+
"""Apply overrides only to attributes exposed by the current config node."""
|
|
22
33
|
for attr, value in kwargs.items():
|
|
23
34
|
if hasattr(obj, attr):
|
|
24
|
-
setattr(obj, attr, value)
|
|
35
|
+
setattr(obj, attr, _normalize_attr_value(obj, attr, value))
|
|
25
36
|
|
|
26
37
|
|
|
27
38
|
def _mutate_common_config_tree(config, visited: set[int] | None = None) -> None:
|
|
39
|
+
"""Shrink nested configs recursively so meta-model construction stays lightweight."""
|
|
28
40
|
if config is None or isinstance(config, (int, float, str, bool, list, tuple, dict)):
|
|
29
41
|
return
|
|
30
42
|
|
|
@@ -125,6 +137,7 @@ def _mutate_common_config_tree(config, visited: set[int] | None = None) -> None:
|
|
|
125
137
|
|
|
126
138
|
|
|
127
139
|
def _build_model_config(case: dict):
|
|
140
|
+
"""Construct a small config tree for one registered public model type."""
|
|
128
141
|
config = _load(case["config_module"], case["config_class"])()
|
|
129
142
|
_mutate_common_config_tree(config)
|
|
130
143
|
|
|
@@ -180,6 +193,7 @@ def _build_model_config(case: dict):
|
|
|
180
193
|
|
|
181
194
|
|
|
182
195
|
def _find_module_hits(model, class_paths: tuple[str, ...]) -> list[tuple[str, str]]:
|
|
196
|
+
"""Collect modules whose class path matches one of the expected targets."""
|
|
183
197
|
hits = []
|
|
184
198
|
wanted = set(class_paths)
|
|
185
199
|
for name, module in model.named_modules():
|
|
@@ -190,16 +204,19 @@ def _find_module_hits(model, class_paths: tuple[str, ...]) -> list[tuple[str, st
|
|
|
190
204
|
|
|
191
205
|
|
|
192
206
|
def _assert_meta_parameters(module) -> None:
|
|
207
|
+
"""Assert that one module keeps all of its parameters on meta."""
|
|
193
208
|
for _, param in module.named_parameters(recurse=True):
|
|
194
209
|
assert param.is_meta
|
|
195
210
|
|
|
196
211
|
|
|
197
212
|
def _assert_all_model_parameters_meta(model) -> None:
|
|
213
|
+
"""Assert that conversion does not materialize weights during meta tests."""
|
|
198
214
|
for _, param in model.named_parameters():
|
|
199
215
|
assert param.is_meta
|
|
200
216
|
|
|
201
217
|
|
|
202
218
|
def _validate_defused_module(case: dict, module) -> None:
|
|
219
|
+
"""Run the case-specific structural checks on one converted module."""
|
|
203
220
|
kind = case["validator"]
|
|
204
221
|
|
|
205
222
|
if kind == "experts":
|
|
@@ -382,13 +399,13 @@ META_MODEL_CASES = [
|
|
|
382
399
|
},
|
|
383
400
|
{
|
|
384
401
|
"model_type": "glm4_moe_lite",
|
|
385
|
-
"mode": "
|
|
402
|
+
"mode": "replace",
|
|
386
403
|
"model_module": "transformers.models.glm4_moe_lite.modeling_glm4_moe_lite",
|
|
387
404
|
"model_class": "Glm4MoeLiteForCausalLM",
|
|
388
405
|
"config_module": "transformers.models.glm4_moe_lite.configuration_glm4_moe_lite",
|
|
389
406
|
"config_class": "Glm4MoeLiteConfig",
|
|
390
|
-
"target_class_paths": ("
|
|
391
|
-
"validator": "
|
|
407
|
+
"target_class_paths": ("defuser.modeling.unfused_moe.glm4_moe_lite.LinearGlm4MoeLiteMoE",),
|
|
408
|
+
"validator": "sparse_block",
|
|
392
409
|
},
|
|
393
410
|
{
|
|
394
411
|
"model_type": "glm4v",
|
|
@@ -729,11 +746,13 @@ META_MODEL_CASES = [
|
|
|
729
746
|
|
|
730
747
|
|
|
731
748
|
def test_meta_model_cases_cover_registered_public_models():
|
|
749
|
+
"""Meta-model coverage should stay aligned with the public registry."""
|
|
732
750
|
assert {case["model_type"] for case in META_MODEL_CASES} == set(MODEL_CONFIG) - {"qwen3_5_moe_text"}
|
|
733
751
|
|
|
734
752
|
|
|
735
753
|
@pytest.mark.parametrize("case", META_MODEL_CASES, ids=[case["model_type"] for case in META_MODEL_CASES])
|
|
736
754
|
def test_each_model_defuses_direct_meta_model(case):
|
|
755
|
+
"""Each registered public model should expose the expected defused modules on meta."""
|
|
737
756
|
if case["mode"] == "replace":
|
|
738
757
|
replace_fused_blocks(case["model_type"])
|
|
739
758
|
|
|
@@ -0,0 +1,62 @@
|
|
|
1
|
+
# SPDX-FileCopyrightText: 2026 ModelCloud.ai
|
|
2
|
+
# SPDX-FileCopyrightText: 2026 qubitium@modelcloud.ai
|
|
3
|
+
# SPDX-License-Identifier: Apache-2.0
|
|
4
|
+
# Contact: qubitium@modelcloud.ai, x.com/qubitium
|
|
5
|
+
|
|
6
|
+
from __future__ import annotations
|
|
7
|
+
|
|
8
|
+
import torch
|
|
9
|
+
import torch.nn as nn
|
|
10
|
+
|
|
11
|
+
from defuser.modeling.moe_experts_interface import _detect_expert_projections, _unfuse_experts_weights_inplace
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
class _UnknownExpertsWithLossProperty(nn.Module):
|
|
15
|
+
"""Regression fixture for expert detection on modules with unrelated properties."""
|
|
16
|
+
|
|
17
|
+
def __init__(self) -> None:
|
|
18
|
+
super().__init__()
|
|
19
|
+
self.expert_weight = nn.Parameter(torch.randn(4, 8, 16))
|
|
20
|
+
self.loss_property_accesses = 0
|
|
21
|
+
|
|
22
|
+
@property
|
|
23
|
+
def loss_function(self):
|
|
24
|
+
self.loss_property_accesses += 1
|
|
25
|
+
raise AssertionError("expert detection should not touch unrelated properties")
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def test_detect_expert_projections_ignores_unrelated_properties() -> None:
|
|
29
|
+
"""Projection detection should ignore unrelated properties on expert fixtures."""
|
|
30
|
+
module = _UnknownExpertsWithLossProperty()
|
|
31
|
+
|
|
32
|
+
detected = _detect_expert_projections(module)
|
|
33
|
+
|
|
34
|
+
assert detected == {
|
|
35
|
+
"expert_weight": {"is_input_proj": True, "output_multiplier": 1},
|
|
36
|
+
}
|
|
37
|
+
assert module.loss_property_accesses == 0
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
class _BufferBackedExperts(nn.Module):
|
|
41
|
+
"""Exercise the buffer-backed fused expert path end to end."""
|
|
42
|
+
|
|
43
|
+
def __init__(self) -> None:
|
|
44
|
+
super().__init__()
|
|
45
|
+
self.register_buffer("gate_up_proj", torch.arange(2 * 6 * 4, dtype=torch.float32).reshape(2, 6, 4))
|
|
46
|
+
self.register_buffer("down_proj", torch.arange(2 * 4 * 3, dtype=torch.float32).reshape(2, 4, 3))
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
def test_unfuse_experts_supports_registered_buffers() -> None:
|
|
50
|
+
"""Buffer-backed fused experts should unfuse into per-expert Linear layers."""
|
|
51
|
+
module = _BufferBackedExperts()
|
|
52
|
+
expected_gate_proj = module.gate_up_proj[0, :3].clone()
|
|
53
|
+
expected_up_proj = module.gate_up_proj[0, 3:].clone()
|
|
54
|
+
expected_down_proj = module.down_proj[0].clone()
|
|
55
|
+
|
|
56
|
+
changed = _unfuse_experts_weights_inplace(module, check_decorator=False)
|
|
57
|
+
|
|
58
|
+
assert changed is True
|
|
59
|
+
expert0 = getattr(module, "0")
|
|
60
|
+
torch.testing.assert_close(expert0.gate_proj.weight, expected_gate_proj)
|
|
61
|
+
torch.testing.assert_close(expert0.up_proj.weight, expected_up_proj)
|
|
62
|
+
torch.testing.assert_close(expert0.down_proj.weight, expected_down_proj)
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
defuser-0.0.17/setup.py
DELETED
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|