@@ -370,14 +370,20 @@ def get_unexpected_file_keys(self, model: Module) -> List[str]:
370
370
# ----- model memory compression/decompression pathways ----- #
371
371
372
372
def compress_model (self , model : Module ):
373
+ """
374
+ Compress a model in memory. Because the model structure is modified in place,
375
+ this method is more memory-efficient than `self.compress`
376
+
377
+ :param model: model containing parameters to compress
378
+ """
373
379
module_to_scheme = map_module_to_scheme (model )
374
380
sparse_compression_targets : Set [str ] = expand_target_names (
375
381
model = model ,
376
382
targets = self .sparsity_config .targets if self .sparsity_config else [],
377
383
ignore = self .sparsity_config .ignore if self .sparsity_config else [],
378
384
)
379
385
380
- for prefix , module in model .named_modules ():
386
+ for prefix , module in tqdm ( model .named_modules (), desc = "Compressing model" ):
381
387
if prefix in module_to_scheme or prefix in sparse_compression_targets :
382
388
state_dict = module .state_dict (prefix = f"{ prefix } ." )
383
389
# quantization first
@@ -409,14 +415,20 @@ def compress_model(self, model: Module):
409
415
module .quantization_status = QuantizationStatus .COMPRESSED
410
416
411
417
def decompress_model (self , model : Module ):
418
+ """
419
+ Decompress a model in memory. Because the model structure is modified in place,
420
+ this method does not require loading some compression parameters from disk
421
+
422
+ :param model: model containing parameters to compress
423
+ """
412
424
module_to_scheme = map_module_to_scheme (model )
413
425
sparse_compression_targets : Set [str ] = expand_target_names (
414
426
model = model ,
415
427
targets = self .sparsity_config .targets if self .sparsity_config else [],
416
428
ignore = self .sparsity_config .ignore if self .sparsity_config else [],
417
429
)
418
430
419
- for prefix , module in model .named_modules ():
431
+ for prefix , module in tqdm ( model .named_modules (), desc = "Decompressing model" ):
420
432
if prefix in module_to_scheme or prefix in sparse_compression_targets :
421
433
state_dict = module .state_dict (prefix = f"{ prefix } ." )
422
434
# sparsity first
0 commit comments