fusion-bench 0.2.22__py3-none-any.whl → 0.2.24__py3-none-any.whl
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.
- fusion_bench/__init__.py +4 -0
- fusion_bench/compat/method/__init__.py +5 -2
- fusion_bench/compat/method/base_algorithm.py +3 -2
- fusion_bench/compat/modelpool/base_pool.py +3 -3
- fusion_bench/compat/taskpool/clip_image_classification.py +1 -1
- fusion_bench/dataset/gpt2_glue.py +1 -1
- fusion_bench/method/__init__.py +12 -2
- fusion_bench/method/analysis/task_vector_cos_similarity.py +95 -12
- fusion_bench/method/analysis/task_vector_violin_plot.py +160 -52
- fusion_bench/method/bitdelta/bitdelta.py +7 -23
- fusion_bench/method/ensemble.py +17 -2
- fusion_bench/method/expert_sparsity/mixtral/dynamic_skipping.py +2 -0
- fusion_bench/method/expert_sparsity/mixtral/layer_wise_pruning.py +2 -0
- fusion_bench/method/expert_sparsity/mixtral/progressive_pruning.py +2 -0
- fusion_bench/method/linear/__init__.py +6 -2
- fusion_bench/method/linear/{simple_average_for_llama.py → simple_average_for_causallm.py} +8 -4
- fusion_bench/method/linear/{task_arithmetic_for_llama.py → task_arithmetic_for_causallm.py} +22 -12
- fusion_bench/method/linear/ties_merging_for_causallm.py +70 -0
- fusion_bench/method/model_stock/__init__.py +1 -0
- fusion_bench/method/model_stock/model_stock.py +309 -0
- fusion_bench/method/regmean/clip_regmean.py +3 -6
- fusion_bench/method/regmean/regmean.py +27 -56
- fusion_bench/method/regmean/utils.py +56 -0
- fusion_bench/method/regmean_plusplus/regmean_plusplus.py +21 -60
- fusion_bench/method/simple_average.py +2 -2
- fusion_bench/method/slerp/__init__.py +1 -1
- fusion_bench/method/slerp/slerp.py +110 -14
- fusion_bench/method/task_arithmetic/task_arithmetic.py +35 -10
- fusion_bench/method/ties_merging/ties_merging.py +22 -6
- fusion_bench/method/we_moe/flan_t5_we_moe.py +9 -20
- fusion_bench/method/wudi/__init__.py +1 -0
- fusion_bench/method/wudi/wudi.py +105 -0
- fusion_bench/mixins/clip_classification.py +26 -6
- fusion_bench/mixins/lightning_fabric.py +4 -0
- fusion_bench/mixins/serialization.py +40 -83
- fusion_bench/modelpool/base_pool.py +1 -1
- fusion_bench/modelpool/causal_lm/causal_lm.py +285 -44
- fusion_bench/modelpool/seq2seq_lm/modelpool.py +146 -0
- fusion_bench/models/hf_clip.py +4 -0
- fusion_bench/models/hf_utils.py +10 -4
- fusion_bench/models/linearized/vision_model.py +6 -6
- fusion_bench/models/model_card_templates/default.md +8 -1
- fusion_bench/models/modeling_smile_mistral/__init__.py +1 -0
- fusion_bench/models/we_moe.py +8 -8
- fusion_bench/models/wrappers/ensemble.py +136 -7
- fusion_bench/scripts/cli.py +2 -2
- fusion_bench/taskpool/base_pool.py +99 -17
- fusion_bench/taskpool/clip_vision/taskpool.py +12 -5
- fusion_bench/taskpool/dummy.py +101 -13
- fusion_bench/taskpool/lm_eval_harness/taskpool.py +80 -0
- fusion_bench/taskpool/nyuv2_taskpool.py +28 -0
- fusion_bench/utils/__init__.py +1 -0
- fusion_bench/utils/data.py +6 -4
- fusion_bench/utils/devices.py +36 -11
- fusion_bench/utils/dtype.py +3 -2
- fusion_bench/utils/lazy_state_dict.py +85 -19
- fusion_bench/utils/packages.py +3 -3
- fusion_bench/utils/parameters.py +0 -2
- fusion_bench/utils/rich_utils.py +7 -3
- fusion_bench/utils/timer.py +92 -10
- {fusion_bench-0.2.22.dist-info → fusion_bench-0.2.24.dist-info}/METADATA +10 -3
- {fusion_bench-0.2.22.dist-info → fusion_bench-0.2.24.dist-info}/RECORD +77 -64
- fusion_bench_config/_get_started/llm_slerp.yaml +12 -0
- fusion_bench_config/method/ensemble/simple_ensemble.yaml +1 -0
- fusion_bench_config/method/linear/{simple_average_for_llama.yaml → simple_average_for_causallm.yaml} +1 -1
- fusion_bench_config/method/linear/task_arithmetic_for_causallm.yaml +4 -0
- fusion_bench_config/method/linear/ties_merging_for_causallm.yaml +13 -0
- fusion_bench_config/method/model_stock/model_stock.yaml +12 -0
- fusion_bench_config/method/slerp/slerp_lm.yaml +4 -0
- fusion_bench_config/method/wudi/wudi.yaml +4 -0
- fusion_bench_config/modelpool/CausalLMPool/{Qwen2.5-1.5B_math_and_coder.yaml → Qwen2.5-1.5B_math_and_code.yaml} +1 -2
- fusion_bench_config/modelpool/CausalLMPool/Qwen2.5-1.5B_three_models.yaml +11 -0
- fusion_bench_config/modelpool/CausalLMPool/llama-7b_3-models_v1.yaml +11 -0
- fusion_bench_config/method/linear/task_arithmetic_for_llama.yaml +0 -4
- {fusion_bench-0.2.22.dist-info → fusion_bench-0.2.24.dist-info}/WHEEL +0 -0
- {fusion_bench-0.2.22.dist-info → fusion_bench-0.2.24.dist-info}/entry_points.txt +0 -0
- {fusion_bench-0.2.22.dist-info → fusion_bench-0.2.24.dist-info}/licenses/LICENSE +0 -0
- {fusion_bench-0.2.22.dist-info → fusion_bench-0.2.24.dist-info}/top_level.txt +0 -0
|
@@ -1,5 +1,5 @@
|
|
|
1
1
|
"""
|
|
2
|
-
Online documentation for this module: https://tanganke.github.io/fusion_bench/modelpool/
|
|
2
|
+
Online documentation for this module: https://tanganke.github.io/fusion_bench/modelpool/llm
|
|
3
3
|
"""
|
|
4
4
|
|
|
5
5
|
import logging
|
|
@@ -26,6 +26,7 @@ from fusion_bench import (
|
|
|
26
26
|
instantiate,
|
|
27
27
|
parse_dtype,
|
|
28
28
|
)
|
|
29
|
+
from fusion_bench.models.hf_utils import create_default_model_card
|
|
29
30
|
from fusion_bench.utils.lazy_state_dict import LazyStateDict
|
|
30
31
|
|
|
31
32
|
log = logging.getLogger(__name__)
|
|
@@ -33,6 +34,44 @@ log = logging.getLogger(__name__)
|
|
|
33
34
|
|
|
34
35
|
@auto_register_config
|
|
35
36
|
class CausalLMPool(BaseModelPool):
|
|
37
|
+
"""A model pool for managing and loading causal language models.
|
|
38
|
+
|
|
39
|
+
This class provides a unified interface for loading and managing multiple
|
|
40
|
+
causal language models, typically used in model fusion and ensemble scenarios.
|
|
41
|
+
It supports both eager and lazy loading strategies, and handles model
|
|
42
|
+
configuration through YAML configs or direct instantiation.
|
|
43
|
+
|
|
44
|
+
The pool can manage models from Hugging Face Hub, local paths, or custom
|
|
45
|
+
configurations. It also provides tokenizer management and model saving
|
|
46
|
+
capabilities with optional Hugging Face Hub integration.
|
|
47
|
+
|
|
48
|
+
Args:
|
|
49
|
+
models: Dictionary or configuration specifying the models to be managed.
|
|
50
|
+
Can contain model names mapped to paths or detailed configurations.
|
|
51
|
+
tokenizer: Tokenizer configuration, either a string path/name or
|
|
52
|
+
a DictConfig with detailed tokenizer settings.
|
|
53
|
+
model_kwargs: Additional keyword arguments passed to model loading.
|
|
54
|
+
Common options include torch_dtype, device_map, etc.
|
|
55
|
+
enable_lazy_loading: Whether to use lazy loading for models. When True,
|
|
56
|
+
models are loaded as LazyStateDict objects instead of actual models,
|
|
57
|
+
which can save memory for large model collections.
|
|
58
|
+
**kwargs: Additional arguments passed to the parent BaseModelPool.
|
|
59
|
+
|
|
60
|
+
Example:
|
|
61
|
+
```python
|
|
62
|
+
>>> pool = CausalLMPool(
|
|
63
|
+
... models={
|
|
64
|
+
... "model_a": "microsoft/DialoGPT-medium",
|
|
65
|
+
... "model_b": "/path/to/local/model"
|
|
66
|
+
... },
|
|
67
|
+
... tokenizer="microsoft/DialoGPT-medium",
|
|
68
|
+
... model_kwargs={"torch_dtype": "bfloat16"}
|
|
69
|
+
... )
|
|
70
|
+
>>> model = pool.load_model("model_a")
|
|
71
|
+
>>> tokenizer = pool.load_tokenizer()
|
|
72
|
+
```
|
|
73
|
+
"""
|
|
74
|
+
|
|
36
75
|
def __init__(
|
|
37
76
|
self,
|
|
38
77
|
models,
|
|
@@ -47,6 +86,20 @@ class CausalLMPool(BaseModelPool):
|
|
|
47
86
|
self.model_kwargs = DictConfig({})
|
|
48
87
|
|
|
49
88
|
def get_model_path(self, model_name: str):
|
|
89
|
+
"""Extract the model path from the model configuration.
|
|
90
|
+
|
|
91
|
+
Args:
|
|
92
|
+
model_name: The name of the model as defined in the models configuration.
|
|
93
|
+
|
|
94
|
+
Returns:
|
|
95
|
+
str: The path or identifier for the model. For string configurations,
|
|
96
|
+
returns the string directly. For dict configurations, extracts
|
|
97
|
+
the 'pretrained_model_name_or_path' field.
|
|
98
|
+
|
|
99
|
+
Raises:
|
|
100
|
+
RuntimeError: If the model configuration is invalid or the model
|
|
101
|
+
name is not found in the configuration.
|
|
102
|
+
"""
|
|
50
103
|
model_name_or_config = self._models[model_name]
|
|
51
104
|
if isinstance(model_name_or_config, str):
|
|
52
105
|
return model_name_or_config
|
|
@@ -56,6 +109,16 @@ class CausalLMPool(BaseModelPool):
|
|
|
56
109
|
raise RuntimeError("Invalid model configuration")
|
|
57
110
|
|
|
58
111
|
def get_model_kwargs(self):
|
|
112
|
+
"""Get processed model keyword arguments for model loading.
|
|
113
|
+
|
|
114
|
+
Converts the stored `model_kwargs` from DictConfig to a regular dictionary
|
|
115
|
+
and processes special arguments like torch_dtype for proper model loading.
|
|
116
|
+
|
|
117
|
+
Returns:
|
|
118
|
+
dict: Processed keyword arguments ready to be passed to model
|
|
119
|
+
loading functions. The torch_dtype field, if present, is
|
|
120
|
+
converted from string to the appropriate torch dtype object.
|
|
121
|
+
"""
|
|
59
122
|
model_kwargs = (
|
|
60
123
|
OmegaConf.to_container(self.model_kwargs, resolve=True)
|
|
61
124
|
if isinstance(self.model_kwargs, DictConfig)
|
|
@@ -71,31 +134,52 @@ class CausalLMPool(BaseModelPool):
|
|
|
71
134
|
model_name_or_config: str | DictConfig,
|
|
72
135
|
*args,
|
|
73
136
|
**kwargs,
|
|
74
|
-
) -> PreTrainedModel:
|
|
75
|
-
"""
|
|
76
|
-
Example of YAML config:
|
|
137
|
+
) -> Union[PreTrainedModel, LazyStateDict]:
|
|
138
|
+
"""Load a causal language model from the model pool.
|
|
77
139
|
|
|
78
|
-
|
|
79
|
-
|
|
80
|
-
|
|
81
|
-
|
|
82
|
-
model_b: path_to_model_b
|
|
83
|
-
```
|
|
140
|
+
This method supports multiple loading strategies:
|
|
141
|
+
1. Loading by model name from the configured model pool
|
|
142
|
+
2. Loading from a direct configuration dictionary
|
|
143
|
+
3. Lazy loading using LazyStateDict for memory efficiency
|
|
84
144
|
|
|
85
|
-
|
|
86
|
-
|
|
87
|
-
|
|
88
|
-
|
|
89
|
-
|
|
90
|
-
|
|
91
|
-
|
|
92
|
-
|
|
93
|
-
|
|
94
|
-
|
|
95
|
-
|
|
96
|
-
|
|
97
|
-
|
|
98
|
-
|
|
145
|
+
The method automatically handles different model configuration formats
|
|
146
|
+
and applies the appropriate loading strategy based on the enable_lazy_loading flag.
|
|
147
|
+
|
|
148
|
+
Args:
|
|
149
|
+
model_name_or_config: Either a string model name that exists in the
|
|
150
|
+
model pool configuration, or a DictConfig/dict containing the
|
|
151
|
+
model configuration directly.
|
|
152
|
+
*args: Additional positional arguments passed to the model constructor.
|
|
153
|
+
**kwargs: Additional keyword arguments passed to the model constructor.
|
|
154
|
+
These will be merged with the pool's model_kwargs.
|
|
155
|
+
|
|
156
|
+
Returns:
|
|
157
|
+
Union[PreTrainedModel, LazyStateDict]: The loaded model. Returns a
|
|
158
|
+
PreTrainedModel for normal loading or a LazyStateDict for lazy loading.
|
|
159
|
+
|
|
160
|
+
Raises:
|
|
161
|
+
RuntimeError: If the model configuration is invalid.
|
|
162
|
+
KeyError: If the model name is not found in the model pool.
|
|
163
|
+
|
|
164
|
+
Example YAML configurations:
|
|
165
|
+
Simple string configuration:
|
|
166
|
+
```yaml
|
|
167
|
+
models:
|
|
168
|
+
_pretrained_: path_to_pretrained_model
|
|
169
|
+
model_a: path_to_model_a
|
|
170
|
+
model_b: path_to_model_b
|
|
171
|
+
```
|
|
172
|
+
|
|
173
|
+
Detailed configuration:
|
|
174
|
+
```yaml
|
|
175
|
+
models:
|
|
176
|
+
_pretrained_:
|
|
177
|
+
_target_: transformers.AutoModelForCausalLM
|
|
178
|
+
pretrained_model_name_or_path: path_to_pretrained_model
|
|
179
|
+
model_a:
|
|
180
|
+
_target_: transformers.AutoModelForCausalLM
|
|
181
|
+
pretrained_model_name_or_path: path_to_model_a
|
|
182
|
+
```
|
|
99
183
|
"""
|
|
100
184
|
model_kwargs = self.get_model_kwargs()
|
|
101
185
|
model_kwargs.update(kwargs)
|
|
@@ -139,23 +223,36 @@ class CausalLMPool(BaseModelPool):
|
|
|
139
223
|
return model
|
|
140
224
|
|
|
141
225
|
def load_tokenizer(self, *args, **kwargs) -> PreTrainedTokenizer:
|
|
142
|
-
"""
|
|
143
|
-
Example of YAML config:
|
|
144
|
-
|
|
145
|
-
```yaml
|
|
146
|
-
tokenizer: google/gemma-2-2b-it # if a plain string, it will be passed to AutoTokenizer.from_pretrained
|
|
147
|
-
```
|
|
226
|
+
"""Load the tokenizer associated with this model pool.
|
|
148
227
|
|
|
149
|
-
|
|
228
|
+
Loads a tokenizer based on the tokenizer configuration provided during
|
|
229
|
+
pool initialization. Supports both simple string paths and detailed
|
|
230
|
+
configuration dictionaries.
|
|
150
231
|
|
|
151
|
-
|
|
152
|
-
|
|
153
|
-
|
|
154
|
-
pretrained_model_name_or_path: google/gemma-2-2b-it
|
|
155
|
-
```
|
|
232
|
+
Args:
|
|
233
|
+
*args: Additional positional arguments passed to the tokenizer constructor.
|
|
234
|
+
**kwargs: Additional keyword arguments passed to the tokenizer constructor.
|
|
156
235
|
|
|
157
236
|
Returns:
|
|
158
|
-
PreTrainedTokenizer: The tokenizer.
|
|
237
|
+
PreTrainedTokenizer: The loaded tokenizer instance.
|
|
238
|
+
|
|
239
|
+
Raises:
|
|
240
|
+
AssertionError: If no tokenizer is defined in the configuration.
|
|
241
|
+
|
|
242
|
+
Example YAML configurations:
|
|
243
|
+
Simple string configuration:
|
|
244
|
+
```yaml
|
|
245
|
+
tokenizer: google/gemma-2-2b-it
|
|
246
|
+
```
|
|
247
|
+
|
|
248
|
+
Detailed configuration:
|
|
249
|
+
```yaml
|
|
250
|
+
tokenizer:
|
|
251
|
+
_target_: transformers.AutoTokenizer
|
|
252
|
+
pretrained_model_name_or_path: google/gemma-2-2b-it
|
|
253
|
+
use_fast: true
|
|
254
|
+
padding_side: left
|
|
255
|
+
```
|
|
159
256
|
"""
|
|
160
257
|
assert self.tokenizer is not None, "Tokenizer is not defined in the config"
|
|
161
258
|
log.info("Loading tokenizer.", stacklevel=2)
|
|
@@ -175,17 +272,54 @@ class CausalLMPool(BaseModelPool):
|
|
|
175
272
|
save_tokenizer: bool = False,
|
|
176
273
|
tokenizer_kwargs=None,
|
|
177
274
|
tokenizer: Optional[PreTrainedTokenizer] = None,
|
|
275
|
+
algorithm_config: Optional[DictConfig] = None,
|
|
276
|
+
description: Optional[str] = None,
|
|
277
|
+
base_model_in_modelcard: bool = True,
|
|
178
278
|
**kwargs,
|
|
179
279
|
):
|
|
180
|
-
"""
|
|
181
|
-
|
|
280
|
+
"""Save a model to the specified path with optional tokenizer and Hub upload.
|
|
281
|
+
|
|
282
|
+
This method provides comprehensive model saving capabilities including
|
|
283
|
+
optional tokenizer saving, dtype conversion, model card creation, and
|
|
284
|
+
Hugging Face Hub upload. The model is saved in the standard Hugging Face format.
|
|
182
285
|
|
|
183
286
|
Args:
|
|
184
|
-
model
|
|
185
|
-
path
|
|
186
|
-
|
|
187
|
-
|
|
188
|
-
|
|
287
|
+
model: The PreTrainedModel instance to be saved.
|
|
288
|
+
path: The local path where the model will be saved. Supports tilde
|
|
289
|
+
expansion for home directory paths.
|
|
290
|
+
push_to_hub: Whether to push the saved model to the Hugging Face Hub.
|
|
291
|
+
Requires proper authentication and repository permissions.
|
|
292
|
+
model_dtype: Optional string specifying the target dtype for the model
|
|
293
|
+
before saving (e.g., "float16", "bfloat16"). The model will be
|
|
294
|
+
converted to this dtype before saving.
|
|
295
|
+
save_tokenizer: Whether to save the tokenizer alongside the model.
|
|
296
|
+
If True, the tokenizer will be loaded using the pool's tokenizer
|
|
297
|
+
configuration and saved to the same path.
|
|
298
|
+
tokenizer_kwargs: Additional keyword arguments for tokenizer loading
|
|
299
|
+
when save_tokenizer is True.
|
|
300
|
+
tokenizer: Optional pre-loaded tokenizer instance. If provided, this
|
|
301
|
+
tokenizer will be saved regardless of the save_tokenizer flag.
|
|
302
|
+
algorithm_config: Optional DictConfig containing algorithm configuration.
|
|
303
|
+
If provided, a model card will be created with algorithm details.
|
|
304
|
+
description: Optional description for the model card. If not provided
|
|
305
|
+
and algorithm_config is given, a default description will be generated.
|
|
306
|
+
**kwargs: Additional keyword arguments passed to the model's
|
|
307
|
+
save_pretrained method.
|
|
308
|
+
|
|
309
|
+
Example:
|
|
310
|
+
```python
|
|
311
|
+
>>> pool = CausalLMPool(models=..., tokenizer=...)
|
|
312
|
+
>>> model = pool.load_model("my_model")
|
|
313
|
+
>>> pool.save_model(
|
|
314
|
+
... model,
|
|
315
|
+
... "/path/to/save",
|
|
316
|
+
... save_tokenizer=True,
|
|
317
|
+
... model_dtype="float16",
|
|
318
|
+
... push_to_hub=True,
|
|
319
|
+
... algorithm_config=algorithm_config,
|
|
320
|
+
... description="Custom merged model"
|
|
321
|
+
... )
|
|
322
|
+
```
|
|
189
323
|
"""
|
|
190
324
|
path = os.path.expanduser(path)
|
|
191
325
|
# NOTE: if tokenizer is provided, it will be saved regardless of `save_tokenizer`
|
|
@@ -207,11 +341,75 @@ class CausalLMPool(BaseModelPool):
|
|
|
207
341
|
**kwargs,
|
|
208
342
|
)
|
|
209
343
|
|
|
344
|
+
# Create and save model card if algorithm_config is provided
|
|
345
|
+
if algorithm_config is not None:
|
|
346
|
+
if description is None:
|
|
347
|
+
description = "Model created using FusionBench."
|
|
348
|
+
model_card_str = create_default_model_card(
|
|
349
|
+
base_model=(
|
|
350
|
+
self.get_model_path("_pretrained_")
|
|
351
|
+
if base_model_in_modelcard and self.has_pretrained
|
|
352
|
+
else None
|
|
353
|
+
),
|
|
354
|
+
models=[self.get_model_path(m) for m in self.model_names],
|
|
355
|
+
description=description,
|
|
356
|
+
algorithm_config=algorithm_config,
|
|
357
|
+
modelpool_config=self.config,
|
|
358
|
+
)
|
|
359
|
+
with open(os.path.join(path, "README.md"), "w") as f:
|
|
360
|
+
f.write(model_card_str)
|
|
361
|
+
|
|
210
362
|
|
|
211
363
|
class CausalLMBackbonePool(CausalLMPool):
|
|
364
|
+
"""A specialized model pool that loads only the transformer backbone layers.
|
|
365
|
+
|
|
366
|
+
This class extends CausalLMPool to provide access to just the transformer
|
|
367
|
+
layers (backbone) of causal language models, excluding the language modeling
|
|
368
|
+
head and embeddings. This is useful for model fusion scenarios where only
|
|
369
|
+
the core transformer layers are needed.
|
|
370
|
+
|
|
371
|
+
The class automatically extracts the `model.layers` component from loaded
|
|
372
|
+
AutoModelForCausalLM instances, providing direct access to the transformer
|
|
373
|
+
blocks. Lazy loading is not supported for this pool type.
|
|
374
|
+
|
|
375
|
+
Note:
|
|
376
|
+
This pool automatically disables lazy loading as it needs to access
|
|
377
|
+
the internal structure of the model to extract the backbone layers.
|
|
378
|
+
|
|
379
|
+
Example:
|
|
380
|
+
```python
|
|
381
|
+
>>> backbone_pool = CausalLMBackbonePool(
|
|
382
|
+
... models={"model_a": "microsoft/DialoGPT-medium"},
|
|
383
|
+
... tokenizer="microsoft/DialoGPT-medium"
|
|
384
|
+
... )
|
|
385
|
+
>>> layers = backbone_pool.load_model("model_a") # Returns nn.ModuleList of transformer layers
|
|
386
|
+
```
|
|
387
|
+
"""
|
|
388
|
+
|
|
212
389
|
def load_model(
|
|
213
390
|
self, model_name_or_config: str | DictConfig, *args, **kwargs
|
|
214
391
|
) -> Module:
|
|
392
|
+
"""Load only the transformer backbone layers from a causal language model.
|
|
393
|
+
|
|
394
|
+
This method loads a complete causal language model and then extracts
|
|
395
|
+
only the transformer layers (backbone), discarding the embedding layers
|
|
396
|
+
and language modeling head. This is useful for model fusion scenarios
|
|
397
|
+
where only the core transformer computation is needed.
|
|
398
|
+
|
|
399
|
+
Args:
|
|
400
|
+
model_name_or_config: Either a string model name from the pool
|
|
401
|
+
configuration or a DictConfig with model loading parameters.
|
|
402
|
+
*args: Additional positional arguments passed to the parent load_model method.
|
|
403
|
+
**kwargs: Additional keyword arguments passed to the parent load_model method.
|
|
404
|
+
|
|
405
|
+
Returns:
|
|
406
|
+
Module: The transformer layers (typically a nn.ModuleList) containing
|
|
407
|
+
the core transformer blocks without embeddings or output heads.
|
|
408
|
+
|
|
409
|
+
Note:
|
|
410
|
+
Lazy loading is automatically disabled for this method as it needs
|
|
411
|
+
to access the internal model structure to extract the layers.
|
|
412
|
+
"""
|
|
215
413
|
if self.enable_lazy_loading:
|
|
216
414
|
log.warning(
|
|
217
415
|
"CausalLMBackbonePool does not support lazy loading. "
|
|
@@ -231,6 +429,49 @@ def load_peft_causal_lm(
|
|
|
231
429
|
is_trainable: bool = True,
|
|
232
430
|
merge_and_unload: bool = False,
|
|
233
431
|
):
|
|
432
|
+
"""Load a causal language model with PEFT (Parameter-Efficient Fine-Tuning) adapters.
|
|
433
|
+
|
|
434
|
+
This function loads a base causal language model and applies PEFT adapters
|
|
435
|
+
(such as LoRA, AdaLoRA, or other parameter-efficient fine-tuning methods)
|
|
436
|
+
to create a fine-tuned model. It supports both keeping the adapters separate
|
|
437
|
+
or merging them into the base model.
|
|
438
|
+
|
|
439
|
+
Args:
|
|
440
|
+
base_model_path: Path or identifier for the base causal language model.
|
|
441
|
+
Can be a Hugging Face model name or local path.
|
|
442
|
+
peft_model_path: Path to the PEFT adapter configuration and weights.
|
|
443
|
+
This should contain the adapter_config.json and adapter weights.
|
|
444
|
+
torch_dtype: The torch data type to use for the model. Common options
|
|
445
|
+
include "float16", "bfloat16", "float32". Defaults to "bfloat16".
|
|
446
|
+
is_trainable: Whether the loaded PEFT model should be trainable.
|
|
447
|
+
Set to False for inference-only usage to save memory.
|
|
448
|
+
merge_and_unload: Whether to merge the PEFT adapters into the base model
|
|
449
|
+
and unload the adapter weights. When True, returns a standard
|
|
450
|
+
PreTrainedModel instead of a PeftModel.
|
|
451
|
+
|
|
452
|
+
Returns:
|
|
453
|
+
Union[PeftModel, PreTrainedModel]: The loaded model with PEFT adapters.
|
|
454
|
+
Returns a PeftModel if merge_and_unload is False, or a PreTrainedModel
|
|
455
|
+
if the adapters are merged and unloaded.
|
|
456
|
+
|
|
457
|
+
Example:
|
|
458
|
+
```python
|
|
459
|
+
>>> # Load model with adapters for training
|
|
460
|
+
>>> model = load_peft_causal_lm(
|
|
461
|
+
... "microsoft/DialoGPT-medium",
|
|
462
|
+
... "/path/to/lora/adapters",
|
|
463
|
+
... is_trainable=True
|
|
464
|
+
... )
|
|
465
|
+
|
|
466
|
+
>>> # Load and merge adapters for inference
|
|
467
|
+
>>> merged_model = load_peft_causal_lm(
|
|
468
|
+
... "microsoft/DialoGPT-medium",
|
|
469
|
+
... "/path/to/lora/adapters",
|
|
470
|
+
... merge_and_unload=True,
|
|
471
|
+
... is_trainable=False
|
|
472
|
+
... )
|
|
473
|
+
```
|
|
474
|
+
"""
|
|
234
475
|
base_model = AutoModelForCausalLM.from_pretrained(
|
|
235
476
|
base_model_path, torch_dtype=torch_dtype
|
|
236
477
|
)
|
|
@@ -18,6 +18,48 @@ def load_lora_model(
|
|
|
18
18
|
is_trainable: bool = True,
|
|
19
19
|
merge_and_unload: bool = True,
|
|
20
20
|
):
|
|
21
|
+
"""Load a sequence-to-sequence model with LoRA (Low-Rank Adaptation) fine-tuning.
|
|
22
|
+
|
|
23
|
+
This function loads a base sequence-to-sequence language model and applies
|
|
24
|
+
LoRA adapters for parameter-efficient fine-tuning. LoRA allows for efficient
|
|
25
|
+
adaptation of large models by adding trainable low-rank matrices to the
|
|
26
|
+
existing weights without modifying the original parameters.
|
|
27
|
+
|
|
28
|
+
Args:
|
|
29
|
+
base_model_path: Path or identifier for the base sequence-to-sequence model.
|
|
30
|
+
Can be a Hugging Face model name (e.g., "t5-base") or local path.
|
|
31
|
+
peft_model_path: Path to the directory containing LoRA adapter weights
|
|
32
|
+
and configuration. Should include adapter_config.json and adapter weights.
|
|
33
|
+
is_trainable: Whether the loaded model should be trainable. Set to False
|
|
34
|
+
for inference-only usage to save memory and computation.
|
|
35
|
+
merge_and_unload: Whether to merge the LoRA weights into the base model
|
|
36
|
+
and unload the adapter. When True, returns a standard model instead
|
|
37
|
+
of a PeftModel, which can be more efficient for inference.
|
|
38
|
+
|
|
39
|
+
Returns:
|
|
40
|
+
Union[PeftModel, AutoModelForSeq2SeqLM]: The loaded model with LoRA
|
|
41
|
+
adapters. Returns a PeftModel if merge_and_unload is False, or
|
|
42
|
+
a standard AutoModelForSeq2SeqLM if adapters are merged.
|
|
43
|
+
|
|
44
|
+
Example:
|
|
45
|
+
```python
|
|
46
|
+
>>> # Load model with separate adapters for training
|
|
47
|
+
>>> model = load_lora_model(
|
|
48
|
+
... "t5-base",
|
|
49
|
+
... "/path/to/lora/adapters",
|
|
50
|
+
... is_trainable=True,
|
|
51
|
+
... merge_and_unload=False
|
|
52
|
+
... )
|
|
53
|
+
|
|
54
|
+
>>> # Load and merge adapters for efficient inference
|
|
55
|
+
>>> merged_model = load_lora_model(
|
|
56
|
+
... "t5-base",
|
|
57
|
+
... "/path/to/lora/adapters",
|
|
58
|
+
... is_trainable=False,
|
|
59
|
+
... merge_and_unload=True
|
|
60
|
+
... )
|
|
61
|
+
```
|
|
62
|
+
"""
|
|
21
63
|
base_model = AutoModelForSeq2SeqLM.from_pretrained(base_model_path)
|
|
22
64
|
model = PeftModel.from_pretrained(
|
|
23
65
|
base_model,
|
|
@@ -30,6 +72,46 @@ def load_lora_model(
|
|
|
30
72
|
|
|
31
73
|
|
|
32
74
|
class Seq2SeqLMPool(BaseModelPool):
|
|
75
|
+
"""A model pool specialized for sequence-to-sequence language models.
|
|
76
|
+
|
|
77
|
+
This model pool provides management and loading capabilities for sequence-to-sequence
|
|
78
|
+
(seq2seq) language models such as T5, BART, and mT5. It extends the base model pool
|
|
79
|
+
functionality with seq2seq-specific features including tokenizer management and
|
|
80
|
+
model configuration handling.
|
|
81
|
+
|
|
82
|
+
Seq2seq models are particularly useful for tasks that require generating output
|
|
83
|
+
sequences from input sequences, such as translation, summarization, question
|
|
84
|
+
answering, and text generation. This pool streamlines the process of loading
|
|
85
|
+
and configuring multiple seq2seq models for fusion and ensemble scenarios.
|
|
86
|
+
|
|
87
|
+
Key Features:
|
|
88
|
+
- Specialized loading for AutoModelForSeq2SeqLM models
|
|
89
|
+
- Integrated tokenizer management
|
|
90
|
+
- Support for model-specific keyword arguments
|
|
91
|
+
- Automatic dtype parsing and configuration
|
|
92
|
+
- Compatible with PEFT (Parameter-Efficient Fine-Tuning) adapters
|
|
93
|
+
|
|
94
|
+
Attributes:
|
|
95
|
+
_tokenizer: Configuration for the tokenizer associated with the models
|
|
96
|
+
_model_kwargs: Default keyword arguments applied to all model loading operations
|
|
97
|
+
|
|
98
|
+
Example:
|
|
99
|
+
```python
|
|
100
|
+
pool = Seq2SeqLMPool(
|
|
101
|
+
models={
|
|
102
|
+
"t5_base": "t5-base",
|
|
103
|
+
"t5_large": "t5-large",
|
|
104
|
+
"custom_model": "/path/to/local/model"
|
|
105
|
+
},
|
|
106
|
+
tokenizer={"_target_": "transformers.T5Tokenizer",
|
|
107
|
+
"pretrained_model_name_or_path": "t5-base"},
|
|
108
|
+
model_kwargs={"torch_dtype": "float16", "device_map": "auto"}
|
|
109
|
+
)
|
|
110
|
+
model = pool.load_model("t5_base")
|
|
111
|
+
tokenizer = pool.load_tokenizer()
|
|
112
|
+
```
|
|
113
|
+
"""
|
|
114
|
+
|
|
33
115
|
_config_mapping = BaseModelPool._config_mapping | {
|
|
34
116
|
"_tokenizer": "tokenizer",
|
|
35
117
|
"_model_kwargs": "model_kwargs",
|
|
@@ -43,6 +125,35 @@ class Seq2SeqLMPool(BaseModelPool):
|
|
|
43
125
|
model_kwargs: Optional[DictConfig] = None,
|
|
44
126
|
**kwargs,
|
|
45
127
|
):
|
|
128
|
+
"""Initialize the sequence-to-sequence language model pool.
|
|
129
|
+
|
|
130
|
+
Sets up the model pool with configurations for models, tokenizer, and
|
|
131
|
+
default model loading parameters. Automatically processes model kwargs
|
|
132
|
+
to handle special configurations like torch_dtype parsing.
|
|
133
|
+
|
|
134
|
+
Args:
|
|
135
|
+
models: Configuration dictionary specifying the seq2seq models to manage.
|
|
136
|
+
Keys are model names, values can be model paths/names or detailed configs.
|
|
137
|
+
tokenizer: Configuration for the tokenizer to use with the models.
|
|
138
|
+
Can be a simple path/name or detailed configuration with _target_.
|
|
139
|
+
model_kwargs: Default keyword arguments applied to all model loading
|
|
140
|
+
operations. Common options include torch_dtype, device_map, etc.
|
|
141
|
+
The torch_dtype field is automatically parsed from string to dtype.
|
|
142
|
+
**kwargs: Additional arguments passed to the parent BaseModelPool.
|
|
143
|
+
|
|
144
|
+
Example:
|
|
145
|
+
```python
|
|
146
|
+
pool = Seq2SeqLMPool(
|
|
147
|
+
models={
|
|
148
|
+
"base": "t5-base",
|
|
149
|
+
"large": {"_target_": "transformers.AutoModelForSeq2SeqLM",
|
|
150
|
+
"pretrained_model_name_or_path": "t5-large"}
|
|
151
|
+
},
|
|
152
|
+
tokenizer="t5-base",
|
|
153
|
+
model_kwargs={"torch_dtype": "bfloat16"}
|
|
154
|
+
)
|
|
155
|
+
```
|
|
156
|
+
"""
|
|
46
157
|
super().__init__(models, **kwargs)
|
|
47
158
|
self._tokenizer = tokenizer
|
|
48
159
|
self._model_kwargs = model_kwargs
|
|
@@ -55,11 +166,46 @@ class Seq2SeqLMPool(BaseModelPool):
|
|
|
55
166
|
)
|
|
56
167
|
|
|
57
168
|
def load_model(self, model_name_or_config: str | DictConfig, *args, **kwargs):
|
|
169
|
+
"""Load a sequence-to-sequence language model from the pool.
|
|
170
|
+
|
|
171
|
+
Loads a seq2seq model using the parent class loading mechanism while
|
|
172
|
+
automatically applying the pool's default model kwargs. The method
|
|
173
|
+
merges the pool's model_kwargs with any additional kwargs provided,
|
|
174
|
+
giving priority to the explicitly provided kwargs.
|
|
175
|
+
|
|
176
|
+
Args:
|
|
177
|
+
model_name_or_config: Either a string model name from the pool
|
|
178
|
+
configuration or a DictConfig containing model loading parameters.
|
|
179
|
+
*args: Additional positional arguments passed to the parent load_model method.
|
|
180
|
+
**kwargs: Additional keyword arguments that override the pool's default
|
|
181
|
+
model_kwargs. Common options include device, torch_dtype, etc.
|
|
182
|
+
|
|
183
|
+
Returns:
|
|
184
|
+
AutoModelForSeq2SeqLM: The loaded sequence-to-sequence language model.
|
|
185
|
+
"""
|
|
58
186
|
model_kwargs = deepcopy(self._model_kwargs)
|
|
59
187
|
model_kwargs.update(kwargs)
|
|
60
188
|
return super().load_model(model_name_or_config, *args, **model_kwargs)
|
|
61
189
|
|
|
62
190
|
def load_tokenizer(self, *args, **kwargs):
|
|
191
|
+
"""Load the tokenizer associated with the sequence-to-sequence models.
|
|
192
|
+
|
|
193
|
+
Loads a tokenizer based on the tokenizer configuration provided during
|
|
194
|
+
pool initialization. The tokenizer should be compatible with the seq2seq
|
|
195
|
+
models in the pool and is typically used for preprocessing input text
|
|
196
|
+
and postprocessing generated output.
|
|
197
|
+
|
|
198
|
+
Args:
|
|
199
|
+
*args: Additional positional arguments passed to the tokenizer constructor.
|
|
200
|
+
**kwargs: Additional keyword arguments passed to the tokenizer constructor.
|
|
201
|
+
|
|
202
|
+
Returns:
|
|
203
|
+
PreTrainedTokenizer: The loaded tokenizer instance compatible with
|
|
204
|
+
the seq2seq models in this pool.
|
|
205
|
+
|
|
206
|
+
Raises:
|
|
207
|
+
AssertionError: If no tokenizer configuration is provided.
|
|
208
|
+
"""
|
|
63
209
|
assert self._tokenizer is not None, "Tokenizer is not defined in the config"
|
|
64
210
|
tokenizer = isinstance(self._tokenizer, *args, **kwargs)
|
|
65
211
|
return tokenizer
|
fusion_bench/models/hf_clip.py
CHANGED
|
@@ -195,5 +195,9 @@ class HFCLIPClassifier(nn.Module):
|
|
|
195
195
|
pass
|
|
196
196
|
elif isinstance(image_embeds, BaseModelOutputWithPooling):
|
|
197
197
|
image_embeds = image_embeds[1]
|
|
198
|
+
elif isinstance(image_embeds, dict) and "pooler_output" in image_embeds:
|
|
199
|
+
image_embeds = image_embeds["pooler_output"]
|
|
200
|
+
else:
|
|
201
|
+
raise ValueError("Unsupported output type from vision model outputs")
|
|
198
202
|
image_embeds = self.clip_model.visual_projection(image_embeds)
|
|
199
203
|
return image_embeds
|
fusion_bench/models/hf_utils.py
CHANGED
|
@@ -42,12 +42,14 @@ def load_model_card_template(basename: str) -> str:
|
|
|
42
42
|
FileNotFoundError: If the template file is not found in any of the search locations.
|
|
43
43
|
"""
|
|
44
44
|
if os.path.exists(basename):
|
|
45
|
-
|
|
45
|
+
with open(basename, "r") as f:
|
|
46
|
+
return f.read()
|
|
46
47
|
|
|
47
48
|
for template_dir in MODEL_CARD_TEMPLATE_DIRS:
|
|
48
49
|
template_path = os.path.join(template_dir, basename)
|
|
49
50
|
if os.path.exists(template_path):
|
|
50
|
-
|
|
51
|
+
with open(template_path, "r") as f:
|
|
52
|
+
return f.read()
|
|
51
53
|
|
|
52
54
|
raise FileNotFoundError(f"Model card template '{basename}' not found.")
|
|
53
55
|
|
|
@@ -141,6 +143,9 @@ def save_pretrained_with_remote_code(
|
|
|
141
143
|
|
|
142
144
|
def create_default_model_card(
|
|
143
145
|
models: list[str],
|
|
146
|
+
base_model: Optional[str] = None,
|
|
147
|
+
title: str = "Deep Model Fusion",
|
|
148
|
+
tags: list[str] = ["fusion-bench", "merge"],
|
|
144
149
|
description=None,
|
|
145
150
|
algorithm_config: DictConfig = None,
|
|
146
151
|
modelpool_config: DictConfig = None,
|
|
@@ -149,10 +154,11 @@ def create_default_model_card(
|
|
|
149
154
|
|
|
150
155
|
template: Template = Template(load_model_card_template("default.md"))
|
|
151
156
|
card = template.render(
|
|
157
|
+
base_model=base_model,
|
|
152
158
|
models=models,
|
|
153
159
|
library_name="transformers",
|
|
154
|
-
|
|
155
|
-
|
|
160
|
+
title=title,
|
|
161
|
+
tags=tags,
|
|
156
162
|
description=description,
|
|
157
163
|
algorithm_config_str=try_to_yaml(algorithm_config),
|
|
158
164
|
modelpool_config_str=try_to_yaml(modelpool_config),
|