diff --git a/docs/source/getting_started.md b/docs/source/getting_started.md index c5c99d6..986801c 100644 --- a/docs/source/getting_started.md +++ b/docs/source/getting_started.md @@ -26,6 +26,25 @@ config.quantization_parameters.default_weight_fractional_bits = 3. config.quantization_parameters.use_relu_multiplier = False ``` +### MDMM pruning (constraint-based) + +MDMM (Modified Differential Method of Multipliers) prunes by enforcing a *constraint* on a chosen sparsity metric instead of adding a fixed penalty. The `metric_type` field picks what is constrained — including two hardware-aware options that target FPGA resources directly. + +```python +from pquant import mdmm_config, mdmm_fpga_config, mdmm_paca_config + +config = mdmm_config() # default: UnstructuredSparsity metric +config = mdmm_fpga_config() # hardware-aware: FPGAAwareSparsity (DSP/BRAM grouping) +config = mdmm_paca_config() # hardware-aware: PACAPatternSparsity (dominant conv-kernel patterns) + +# switch the constrained metric and its parameters +config.pruning_parameters.metric_type = "FPGAAwareSparsity" +config.pruning_parameters.target_resource = "DSP" # or "BRAM" +config.pruning_parameters.rf = 4 +``` + +Training proceeds in three phases handled by the training loop: a warm-up where the constraint is inactive, an active phase where the constraint loss is applied and the prune mask is tracked, and fine-tuning where the mask is frozen and applied. The available `metric_type` options and their parameters are listed in the [Usage Reference](reference.md). + ### Building a model PQuantML supports two ways of defining compressed models. Below we illustrate both approaches using a simple jet-tagging architecture. diff --git a/docs/source/reference.md b/docs/source/reference.md index 10802c7..7cc881a 100644 --- a/docs/source/reference.md +++ b/docs/source/reference.md @@ -169,7 +169,7 @@ There are more details about every pruning method: | `pruning_method` | str | `mdmm` | Selects this pruning schema. | | `constraint_type` | ConstraintType | `"Equality"` | Constraint form: equality / ≤ / ≥. | | `target_value` | float | `0.0` | Target value for the chosen metric. | -| `metric_type` | MetricType | `"UnstructuredSparsity"` | Specifies which metric is constrained. | +| `metric_type` | MetricType | `"UnstructuredSparsity"` | Quantity the constraint acts on — see **MDMM metric types** below. | | `target_sparsity` | float | `0.9` | Target sparsity when constraining sparsity. | | `rf` | int | `1` | Regularization / frequency parameter. | | `epsilon` | float | `1.0e-03` | Feasibility tolerance. | @@ -180,6 +180,43 @@ There are more details about every pruning method: | `scale_mode` | `"mean"` \| `"sum"` | `"mean"` | Aggregation mode for penalties. | +##### MDMM metric types + +The `metric_type` field selects which quantity the MDMM constraint drives. The first two are magnitude-based; the last two are hardware-aware and act on 4D convolution kernels. + +| **metric_type** | **Constrains** | +|------------------------|-------------------------------------------------------------------------------| +| `UnstructuredSparsity` | Element-wise (L0/L1) sparsity toward `target_sparsity`. | +| `StructuredSparsity` | Fraction of all-zero weight groups of size `rf`. | +| `FPGAAwareSparsity` | Fraction of zero DSP/BRAM weight groups, modelling FPGA resource packing. | +| `PACAPatternSparsity` | Mean distance of each conv kernel to a small set of dominant binary patterns. | + +**`FPGAAwareSparsity` parameters** (used only when `metric_type: FPGAAwareSparsity`): + +| **Field** | **Type** | **Default** | **Description** | +|-------------------|---------------------|-------------|-----------------------------------------------------------------------------| +| `precision` | int | `16` | Weight bit-width used to derive BRAM packing. | +| `target_resource` | `"DSP"` \| `"BRAM"` | `"DSP"` | Hardware resource whose group-sparsity is measured. | +| `bram_width` | int | `36` | BRAM word width; sets how many DSP groups pack into one BRAM (`BRAM` only). | + +Weights are grouped into DSP blocks of size `rf`; for `target_resource: BRAM`, `c = bram_width // precision` (or `2*bram_width // precision` when not divisible) consecutive DSP groups pack into one BRAM block. The metric reports the fraction of such groups whose L2 norm is below `epsilon`. + +**`PACAPatternSparsity` parameters** (used only when `metric_type: PACAPatternSparsity`): + +| **Field** | **Type** | **Default** | **Description** | +|------------------------|-------------------------------------------------|--------------------|--------------------------------------------------------------| +| `num_patterns_to_keep` | int | `16` | Maximum number of dominant kernel patterns retained. | +| `beta` | float | `0.75` | Cumulative pattern-frequency coverage kept, in `[0, 1]`. | +| `distance_metric` | `"hamming"` \| `"valued_hamming"` \| `"cosine"` | `"valued_hamming"` | Distance from each kernel to its closest dominant pattern. | + +```{note} +`PACAPatternSparsity` always pairs with an equality constraint at target `0` (driving every kernel onto a dominant pattern); the config model sets `constraint_type` and `target_value` for you. During fine-tuning the kernels are projected onto their closest dominant pattern. +``` + +```{note} +The hardware-aware metrics operate on 4D convolution weights; for non-convolutional layers `PACAPatternSparsity` is a no-op. Ready-made configs are available via `mdmm_fpga_config()` and `mdmm_paca_config()`. +``` + Optionally, there is also FITCompress method implemented for PyTorch: ### FitCompress method | **Field** | **Type** | **Default** | **Description** | diff --git a/src/pquant/__init__.py b/src/pquant/__init__.py index 6c27f9d..e2878c3 100644 --- a/src/pquant/__init__.py +++ b/src/pquant/__init__.py @@ -16,6 +16,8 @@ load_from_dictionary, load_from_file, mdmm_config, + mdmm_fpga_config, + mdmm_paca_config, pdp_config, wanda_config, ) @@ -58,6 +60,8 @@ _forwards.append("cs_config") _forwards.append("dst_config") _forwards.append("mdmm_config") + _forwards.append("mdmm_fpga_config") + _forwards.append("mdmm_paca_config") _forwards.append("pdp_config") _forwards.append("wanda_config") _forwards.append("fitcompress_config") @@ -81,6 +85,8 @@ load_from_dictionary, load_from_file, mdmm_config, + mdmm_fpga_config, + mdmm_paca_config, pdp_config, wanda_config, ) @@ -114,6 +120,8 @@ _forwards.append("cs_config") _forwards.append("dst_config") _forwards.append("mdmm_config") + _forwards.append("mdmm_fpga_config") + _forwards.append("mdmm_paca_config") _forwards.append("pdp_config") _forwards.append("wanda_config") _forwards.append("load_from_file") diff --git a/src/pquant/configs/config_mdmm_fpga.yaml b/src/pquant/configs/config_mdmm_fpga.yaml new file mode 100644 index 0000000..a157d95 --- /dev/null +++ b/src/pquant/configs/config_mdmm_fpga.yaml @@ -0,0 +1,72 @@ +# file: /src/pquant/configs/config_mdmm_fpga.yaml +# MDMM pruning with the hardware-aware FPGA metric (DSP/BRAM resource grouping). + +pruning_parameters: + pruning_method: mdmm + enable_pruning: true + disable_pruning_for_layers: [] # Disable pruning for these layers, even if enable_pruning is true + constraint_type: "Equality" + target_value: 0.0 + metric_type: "FPGAAwareSparsity" + target_sparsity: 0.9 + rf: 4 + epsilon: 1.0e-03 + scale: 50.0 + damping: 1.0 + use_grad: false + constraint_lr: 1.0e-3 + # FPGAAwareSparsityMetric-specific + precision: 16 # weight bit-width + target_resource: "DSP" # "DSP" or "BRAM" + bram_width: 36 # BRAM word width (only used when target_resource == "BRAM") + +quantization_parameters: + enable_quantization: true + default_weight_keep_negatives: 1. + default_weight_integer_bits: 0. + default_weight_fractional_bits: 7. + default_data_keep_negatives: 0. + default_data_integer_bits: 0. + default_data_fractional_bits: 7. + dynamic_data_quantization: false + granularity: "per_tensor" + quantize_input: true + quantize_output: false + hgq_beta: 1e-5 + hgq_gamma: 0.0003 + hgq_heterogeneous: True + layer_specific: {} + use_high_granularity_quantization: false + use_real_tanh: false + use_relu_multiplier: false + use_symmetric_quantization: false + overflow_mode_parameters: SAT + overflow_mode_data: SAT + round_mode: RND +training_parameters: + epochs: 200 + fine_tuning_epochs: 30 + pretraining_epochs: 0 + pruning_first: true + rewind: never + rounds: 1 + save_weights_epoch: -1 +fitcompress_parameters: + enable_fitcompress : false + optimize_quantization : true + quantization_schedule : [7.,4.,3.,2] + pruning_schedule : {start : 0, end : -3, steps : 40} + compression_goal : 0.10 + optimize_pruning : false + greedy_astar : true + approximate : true + f_lambda : 1 +hpo_parameters: + experiment_name: experiment_name + model_name: jet_tagger + num_trials: 1 + sampler: + type: RandomSampler + hyperparameter_search: + numerical: {} + categorical: {} diff --git a/src/pquant/configs/config_mdmm_paca.yaml b/src/pquant/configs/config_mdmm_paca.yaml new file mode 100644 index 0000000..a76a383 --- /dev/null +++ b/src/pquant/configs/config_mdmm_paca.yaml @@ -0,0 +1,74 @@ +# file: /src/pquant/configs/config_mdmm_paca.yaml +# MDMM pruning with the PACA pattern metric. PACA always pairs with an equality +# constraint at target 0 (enforced by the config model); the values below are listed +# for clarity. + +pruning_parameters: + pruning_method: mdmm + enable_pruning: true + disable_pruning_for_layers: [] # Disable pruning for these layers, even if enable_pruning is true + constraint_type: "Equality" # Forced to Equality/0 for PACA by MDMMPruningModel. + target_value: 0.0 + metric_type: "PACAPatternSparsity" + target_sparsity: 0.9 + rf: 1 + epsilon: 1.0e-03 + scale: 1.0 + damping: 1.0 + use_grad: false + constraint_lr: 1.0e-3 + # PACAPatternMetric-specific + num_patterns_to_keep: 16 + beta: 0.85 # cumulative pattern-frequency coverage kept (0..1) + distance_metric: "cosine" # "hamming", "valued_hamming", or "cosine" + +quantization_parameters: + enable_quantization: true + default_weight_keep_negatives: 1. + default_weight_integer_bits: 0. + default_weight_fractional_bits: 7. + default_data_keep_negatives: 0. + default_data_integer_bits: 0. + default_data_fractional_bits: 7. + dynamic_data_quantization: false + granularity: "per_tensor" + quantize_input: true + quantize_output: false + hgq_beta: 1e-5 + hgq_gamma: 0.0003 + hgq_heterogeneous: True + layer_specific: {} + use_high_granularity_quantization: false + use_real_tanh: false + use_relu_multiplier: false + use_symmetric_quantization: false + overflow_mode_parameters: SAT + overflow_mode_data: SAT + round_mode: RND +training_parameters: + epochs: 200 + fine_tuning_epochs: 30 + pretraining_epochs: 0 + pruning_first: true + rewind: never + rounds: 1 + save_weights_epoch: -1 +fitcompress_parameters: + enable_fitcompress : false + optimize_quantization : true + quantization_schedule : [7.,4.,3.,2] + pruning_schedule : {start : 0, end : -3, steps : 40} + compression_goal : 0.10 + optimize_pruning : false + greedy_astar : true + approximate : true + f_lambda : 1 +hpo_parameters: + experiment_name: experiment_name + model_name: jet_tagger + num_trials: 1 + sampler: + type: RandomSampler + hyperparameter_search: + numerical: {} + categorical: {} diff --git a/src/pquant/core/constants.py b/src/pquant/core/constants.py index 6714d98..0fe10c7 100644 --- a/src/pquant/core/constants.py +++ b/src/pquant/core/constants.py @@ -44,3 +44,20 @@ CONFIG_FILE = "config.yaml" N_JOBS = 1 + +# --- Hardware-aware pruning metric constants --- +# Conv-kernel layout -> axis index, used to canonicalise weight layouts in the PACA +# pattern utilities. Keras conv weights are HWIO, Torch conv weights are OIHW. +CONV_LAYOUT_AXES = {"H": 0, "W": 1, "I": 2, "O": 3} +CANONICAL_CONV_LAYOUT = "OIHW" + +# PACAPatternMetric pattern-distance metrics +DISTANCE_HAMMING = "hamming" +DISTANCE_VALUED_HAMMING = "valued_hamming" +DISTANCE_COSINE = "cosine" +PACA_DISTANCE_METRICS = (DISTANCE_HAMMING, DISTANCE_VALUED_HAMMING, DISTANCE_COSINE) + +# FPGAAwareSparsityMetric target hardware resources +TARGET_RESOURCE_DSP = "DSP" +TARGET_RESOURCE_BRAM = "BRAM" +FPGA_TARGET_RESOURCES = (TARGET_RESOURCE_DSP, TARGET_RESOURCE_BRAM) diff --git a/src/pquant/core/hyperparameter_optimization.py b/src/pquant/core/hyperparameter_optimization.py index c9d4289..80e0b99 100644 --- a/src/pquant/core/hyperparameter_optimization.py +++ b/src/pquant/core/hyperparameter_optimization.py @@ -450,6 +450,20 @@ def mdmm_config(): return PQConfig.load_from_file(path) +def mdmm_fpga_config(): + yaml_name = "config_mdmm_fpga.yaml" + parent = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + path = os.path.join(parent, "configs", yaml_name) + return PQConfig.load_from_file(path) + + +def mdmm_paca_config(): + yaml_name = "config_mdmm_paca.yaml" + parent = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + path = os.path.join(parent, "configs", yaml_name) + return PQConfig.load_from_file(path) + + def pdp_config(): yaml_name = "config_pdp.yaml" parent = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) diff --git a/src/pquant/core/keras/pruning_methods/mdmm.py b/src/pquant/core/keras/pruning_methods/mdmm.py index 62b5e01..9a270dc 100644 --- a/src/pquant/core/keras/pruning_methods/mdmm.py +++ b/src/pquant/core/keras/pruning_methods/mdmm.py @@ -14,6 +14,8 @@ LessThanOrEqualConstraint, ) from pquant.core.keras.pruning_methods.metric_functions import ( + FPGAAwareSparsityMetric, + PACAPatternMetric, StructuredSparsityMetric, UnstructuredSparsityMetric, ) @@ -21,6 +23,8 @@ METRIC_REGISTRY = { "UnstructuredSparsity": UnstructuredSparsityMetric, "StructuredSparsity": StructuredSparsityMetric, + "FPGAAwareSparsity": FPGAAwareSparsityMetric, + "PACAPatternSparsity": PACAPatternMetric, } CONSTRAINT_REGISTRY = { @@ -67,6 +71,14 @@ def build(self, input_shape): "l0_mode": l0_mode, "scale_mode": scale_mode, "rf": pruning_parameters.rf, + # FPGAAwareSparsityMetric + "precision": pruning_parameters.precision, + "target_resource": pruning_parameters.target_resource, + "bram_width": pruning_parameters.bram_width, + # PACAPatternMetric + "num_patterns_to_keep": pruning_parameters.num_patterns_to_keep, + "beta": pruning_parameters.beta, + "distance_metric": pruning_parameters.distance_metric, } metric_cls = METRIC_REGISTRY.get(metric_type) @@ -110,9 +122,18 @@ def build(self, input_shape): self.constraint_layer.build(input_shape) super().build(input_shape) + def _compute_hard_mask(self, weight, epsilon): + # During fine-tuning, a metric that defines its own projection (e.g. PACA pattern + # pruning) supplies the mask; otherwise use the magnitude threshold. The layer only + # checks for the capability, so it stays metric-agnostic (no metric-type branching). + metric_fn = getattr(self.constraint_layer, "metric_fn", None) + if self._is_finetuning and hasattr(metric_fn, "get_projection_mask"): + return ops.cast(metric_fn.get_projection_mask(weight), weight.dtype) + return ops.cast(ops.abs(weight) > epsilon, weight.dtype) + def call(self, weight): epsilon = self.config.pruning_parameters.epsilon - hard_mask = ops.cast(ops.abs(weight) > epsilon, weight.dtype) + hard_mask = self._compute_hard_mask(weight, epsilon) not_active = ops.logical_or(self.is_pretraining, self.is_finetuning) self.mask.assign(ops.where(not_active, ops.convert_to_tensor(self.mask), hard_mask)) @@ -131,7 +152,10 @@ def get_hard_mask(self, weight=None): return ops.cast(ops.abs(weight) > epsilon, weight.dtype) def get_layer_sparsity(self, weight): - return ops.sum(self.get_hard_mask(weight)) / ops.size(weight) + # Cast size to the mask dtype: ops.sum(mask) is float but ops.size is int, and the + # TensorFlow backend rejects float/int division (the original float32/int32 bug). + mask = self.get_hard_mask(weight) + return ops.sum(mask) / ops.cast(ops.size(weight), mask.dtype) def calculate_additional_loss(self): # Loss is added via self.add_loss() in call() for model.fit. diff --git a/src/pquant/core/keras/pruning_methods/metric_functions.py b/src/pquant/core/keras/pruning_methods/metric_functions.py index 0f22b5e..6230f46 100644 --- a/src/pquant/core/keras/pruning_methods/metric_functions.py +++ b/src/pquant/core/keras/pruning_methods/metric_functions.py @@ -1,5 +1,13 @@ from keras import ops +from pquant.core.constants import ( + DISTANCE_VALUED_HAMMING, + FPGA_TARGET_RESOURCES, + TARGET_RESOURCE_BRAM, + TARGET_RESOURCE_DSP, +) +from pquant.core.keras.pruning_methods import patterns + class UnstructuredSparsityMetric: """L0-L1 based metric""" @@ -83,3 +91,122 @@ def __call__(self, weight): num_groups = ops.cast(ops.size(group_norms), "float32") return ops.sum(ops.cast(zero_groups, "float32")) / num_groups + + +class FPGAAwareSparsityMetric: + """Hardware-aware sparsity metric for FPGA targets. + + Models how weights are packed into DSP blocks (groups of size ``rf``) and further + into BRAM blocks (groups of ``c`` DSP blocks, where ``c`` derives from ``bram_width`` + and ``precision``). Returns the fraction of zero-valued groups at the chosen + ``target_resource`` level. Constructor inputs are validated by the Pydantic config + model (MDMMPruningModel), so no input asserts are kept here. + """ + + def __init__(self, rf=1, precision=16, target_resource=TARGET_RESOURCE_DSP, bram_width=36, epsilon=1e-3): + self.rf = rf + self.precision = precision + self.target_resource = target_resource + self.bram_width = bram_width + self.epsilon = epsilon + self.c = self._calculate_c() + + def _calculate_c(self): + """Number of consecutive DSP groups packed into a single BRAM block.""" + if self.bram_width % self.precision == 0: + return self.bram_width // self.precision + return (2 * self.bram_width) // self.precision + + def _prepare_weights(self, weight): + """Flatten to (rows, -1) and right-pad the row length to a multiple of rf.""" + original_shape = weight.shape + # 1D (e.g. bias) -> single row; 2D -> as-is; >2D (e.g. Conv2D) -> flatten trailing dims. + if len(original_shape) == 1: + weight_reshaped = ops.reshape(weight, (1, -1)) + elif len(original_shape) > 2: + weight_reshaped = ops.reshape(weight, (original_shape[0], -1)) + else: + weight_reshaped = weight + num_weights = ops.shape(weight_reshaped)[1] + padding_needed = (self.rf - num_weights % self.rf) % self.rf + return ops.pad(weight_reshaped, [[0, 0], [0, padding_needed]]) + + def __call__(self, weight): + prepared = self._prepare_weights(weight) + dsp_groups = ops.reshape(prepared, (prepared.shape[0], -1, self.rf)) + if self.target_resource == TARGET_RESOURCE_DSP: + return self._dsp_sparsity(dsp_groups) + if self.target_resource == TARGET_RESOURCE_BRAM: + return self._bram_sparsity(dsp_groups) + raise ValueError(f"target_resource must be one of {FPGA_TARGET_RESOURCES}, got {self.target_resource!r}") + + def _dsp_sparsity(self, dsp_groups): + """A DSP block is pruned when the L2-norm of its weight group is below epsilon.""" + group_norms = ops.sqrt(ops.sum(ops.square(dsp_groups), axis=-1)) + zero_groups = ops.less_equal(group_norms, self.epsilon) + num_groups = ops.cast(ops.size(group_norms), dsp_groups.dtype) + return ops.sum(ops.cast(zero_groups, dsp_groups.dtype)) / num_groups + + def _bram_sparsity(self, dsp_groups): + """A BRAM block is pruned when the L2-norm of all weights stored in it is below epsilon.""" + if self.c < 1: + raise ValueError( + f"BRAM packing needs precision <= 2*bram_width (got precision={self.precision}, " + f"bram_width={self.bram_width} -> c={self.c})." + ) + num_dsp_groups = ops.shape(dsp_groups)[1] + bram_padding = (self.c - num_dsp_groups % self.c) % self.c + dsp_padded = ops.pad(dsp_groups, [[0, 0], [0, bram_padding], [0, 0]]) + bram_groups = ops.reshape(dsp_padded, (dsp_groups.shape[0], -1, self.c, self.rf)) + bram_norms = ops.sqrt(ops.sum(ops.square(bram_groups), axis=(-1, -2))) + zero_bram = ops.less_equal(bram_norms, self.epsilon) + num_bram = ops.cast(ops.size(bram_norms), dsp_groups.dtype) + return ops.sum(ops.cast(zero_bram, dsp_groups.dtype)) / num_bram + + +class PACAPatternMetric: + """Pattern-based pruning metric (PACA). + + On the first call it selects a small set of dominant binary patterns over the conv + kernels (cached for the metric instance lifetime), then returns the mean distance of + every kernel to its closest dominant pattern. Operates on 4D conv weights only (Keras + kernels are HWIO); returns 0 for non-4D inputs. Exposes ``get_projection_mask`` so the + MDMM layer can snap weights onto their patterns during fine-tuning. Constructor inputs + are validated by the Pydantic config model (MDMMPruningModel). + """ + + def __init__( + self, num_patterns_to_keep=16, beta=0.75, epsilon=1e-5, distance_metric=DISTANCE_VALUED_HAMMING, src="HWIO" + ): + self.num_patterns_to_keep = num_patterns_to_keep + self.beta = beta + self.epsilon = epsilon + self.distance_metric = distance_metric + self.src = src # Keras conv kernels are HWIO + self.dominant_patterns = None + self.valid_mask = None + + def _ensure_patterns(self, weight): + if self.dominant_patterns is None: + _, pats, _ = patterns.kernels_and_patterns(weight, self.src, self.epsilon) + self.dominant_patterns, self.valid_mask = patterns.select_dominant_patterns( + pats, self.num_patterns_to_keep, self.beta + ) + + def __call__(self, weight): + if len(weight.shape) != 4: + return ops.convert_to_tensor(0.0, dtype=weight.dtype) + self._ensure_patterns(weight) + _, distances = patterns.pattern_distances( + weight, self.dominant_patterns, self.valid_mask, self.src, self.epsilon, self.distance_metric + ) + return ops.mean(ops.min(distances, axis=1)) + + def get_projection_mask(self, weight): + # Identity mask if patterns were never selected (e.g. metric never invoked) so the + # MDMM layer's `weight * mask` is a safe no-op. + if self.dominant_patterns is None: + return ops.ones_like(weight) + return patterns.projection_mask( + weight, self.dominant_patterns, self.valid_mask, self.src, self.epsilon, self.distance_metric + ) diff --git a/src/pquant/core/keras/pruning_methods/patterns.py b/src/pquant/core/keras/pruning_methods/patterns.py new file mode 100644 index 0000000..a0888c4 --- /dev/null +++ b/src/pquant/core/keras/pruning_methods/patterns.py @@ -0,0 +1,167 @@ +# @Author: Arghya Ranjan Das +# PACA pattern utilities (Keras backend). +# +# Selects a small set of dominant binary "patterns" (the support, i.e. the non-zero +# positions, of each conv kernel) and measures how far each kernel is from its closest +# dominant pattern. Everything stays in keras.ops (on-device): unique-pattern counting +# is done via bit-pack -> argsort -> bincount over sorted runs, and the dominant set is +# returned at a fixed size (num_patterns_to_keep) together with a validity mask, so no +# host<->device sync or NumPy round-trip of the weight is needed. +# +# Conv weights are canonicalised to OIHW (C_out, C_in, kH, kW) before any pattern logic, +# so both the Keras (HWIO kernel) and Torch (OIHW) layouts are handled by passing `src`. + +import keras +from keras import ops + +from pquant.core.constants import ( + CANONICAL_CONV_LAYOUT, + CONV_LAYOUT_AXES, + DISTANCE_COSINE, + DISTANCE_HAMMING, + DISTANCE_VALUED_HAMMING, +) + +_INF = 1e30 + + +def _layout_to_axes(layout): + if len(layout) != 4 or set(layout) != set("HWIO"): + raise ValueError(f"layout must be a permutation of 'HWIO', got {layout!r}") + return tuple(CONV_LAYOUT_AXES[ch] for ch in layout) + + +def _perm(src, dst): + """Permutation tuple that reorders axes from `src` layout to `dst` layout.""" + s = _layout_to_axes(src) + d = _layout_to_axes(dst) + return tuple(s.index(ax) for ax in d) + + +def convert_conv_layout(w, src, dst=CANONICAL_CONV_LAYOUT): + """Transpose a 4D conv weight from `src` to `dst` layout (no-op if already equal).""" + if src == dst: + return w + perm = _perm(src, dst) + if perm == (0, 1, 2, 3): + return w + return ops.transpose(w, perm) + + +def kernels_and_patterns(w, src, epsilon): + """Flatten a 4D conv weight to per-kernel rows and their binary support. + + Returns (kernels, patterns, (C_out, C_in, kH, kW)): + kernels: (C_out*C_in, kH*kW) float - flattened kernels, canonical OIHW order. + patterns: (C_out*C_in, kH*kW) uint8 - binary support, |w| > epsilon. + """ + w_oihw = convert_conv_layout(w, src=src, dst=CANONICAL_CONV_LAYOUT) + c_out, c_in, kh, kw = w_oihw.shape + kernels = ops.reshape(w_oihw, (c_out * c_in, kh * kw)) + patterns = ops.cast(ops.greater(ops.abs(kernels), epsilon), "uint8") + return kernels, patterns, (c_out, c_in, kh, kw) + + +def _pattern_codes(patterns): + """Bit-pack each binary pattern row into a unique int64 code. + + Two rows share a code iff identical. Valid for kH*kW <= 62 (all realistic conv + kernels); larger kernels are not expected for hardware-aware pattern pruning. + """ + k = patterns.shape[1] + weights = ops.power(ops.full((k,), 2, dtype="int64"), ops.arange(k, dtype="int64")) + return ops.sum(ops.cast(patterns, "int64") * weights, axis=1) # (M,) + + +def select_dominant_patterns(patterns, num_patterns_to_keep, beta): + """Select the most frequent distinct patterns covering `beta` of the total mass. + + Pure keras.ops. Returns (dominant, valid): + dominant: (num_patterns_to_keep, kH*kW) uint8 - the candidate patterns. + valid: (num_patterns_to_keep,) bool - which rows are real selections + (rows past the beta cut / past the number of distinct patterns are + padding and must be ignored downstream). + """ + alpha = int(num_patterns_to_keep) + m, k = patterns.shape[0], patterns.shape[1] + if m == 0: + return ops.zeros((alpha, k), "uint8"), ops.zeros((alpha,), "bool") + + codes = _pattern_codes(patterns) # (M,) int64 + order = ops.argsort(codes) + codes_sorted = ops.take(codes, order) + pat_sorted = ops.take(patterns, order, axis=0) # identical patterns now contiguous + + # First position of each distinct code in sorted order. + not_equal_prev = ops.not_equal(codes_sorted[1:], codes_sorted[:-1]) + is_start = ops.concatenate([ops.ones((1,), "bool"), not_equal_prev], axis=0) # (M,) + + group = ops.cumsum(ops.cast(is_start, "int32")) - 1 # dense group ids 0..U-1 (M,) + group_counts = ops.bincount(group, minlength=m) # counts per group id + count_per_pos = ops.take(group_counts, group) # (M,) + + total = ops.cast(m, "float32") + pdf = ops.where(is_start, ops.cast(count_per_pos, "float32") / total, ops.zeros((m,), "float32")) + + # Order representatives by descending frequency; zero-pdf duplicates sink to the end. + order2 = ops.argsort(-pdf) + pdf_desc = ops.take(pdf, order2) + pat_desc = ops.take(pat_sorted, order2, axis=0) + cdf = ops.cumsum(pdf_desc) + + # keep = min( #patterns to reach beta coverage, alpha cap, #distinct patterns ) + reaches = ops.cast(cdf >= beta, "int32") + has_hit = ops.sum(reaches) > 0 + n_beta = ops.cast(ops.argmax(reaches) + 1, "int32") + n_distinct = ops.sum(ops.cast(is_start, "int32")) + keep = ops.where(has_hit, n_beta, n_distinct) + keep = ops.minimum(ops.minimum(keep, n_distinct), alpha) # 0-d int tensor, <= alpha + + # Static-size top-alpha slice (alpha is a Python int), zero-padded if M < alpha. + pat_top = pat_desc[:alpha] + pad_rows = alpha - int(pat_top.shape[0]) + if pad_rows > 0: + pat_top = ops.concatenate([pat_top, ops.zeros((pad_rows, k), pat_top.dtype)], axis=0) + valid = ops.arange(alpha) < ops.cast(keep, "int32") # (alpha,) bool, stays on-device + return pat_top, valid + + +def _kernel_pattern_distances(kernel_patterns, kernels, dominant_patterns, distance_metric): + """Distance from every kernel to every dominant pattern. -> (M_kernels, alpha).""" + tk = ops.cast(kernel_patterns, kernels.dtype) # (M, K) binary support of kernels + p = ops.cast(dominant_patterns, kernels.dtype) # (alpha, K) + tk_e = ops.expand_dims(tk, 1) # (M, 1, K) + k_e = ops.expand_dims(kernels, 1) # (M, 1, K) + p_e = ops.expand_dims(p, 0) # (1, alpha, K) + + if distance_metric == DISTANCE_HAMMING: + return ops.sum(ops.abs(tk_e - p_e), axis=-1) + if distance_metric == DISTANCE_VALUED_HAMMING: + return ops.sum(ops.abs(tk_e - p_e) * ops.abs(k_e), axis=-1) + if distance_metric == DISTANCE_COSINE: + projected = k_e * p_e + dot = ops.sum(k_e * projected, axis=-1) + denom = ops.norm(k_e, axis=-1) * ops.norm(projected, axis=-1) + keras.backend.epsilon() + return 1.0 - dot / denom + raise ValueError(f"Unsupported distance metric: {distance_metric!r}") + + +def pattern_distances(w, dominant_patterns, valid_mask, src, epsilon, distance_metric): + """Per-kernel distance to each dominant pattern, with invalid patterns masked to +inf.""" + kernels, kernel_patterns, _ = kernels_and_patterns(w, src, epsilon) + distances = _kernel_pattern_distances(kernel_patterns, kernels, dominant_patterns, distance_metric) + distances = ops.where(valid_mask[None, :], distances, ops.cast(_INF, distances.dtype)) + return kernels, distances + + +def projection_mask(w, dominant_patterns, valid_mask, src, epsilon, distance_metric): + """Binary mask (same layout as `w`) projecting each kernel onto its closest dominant pattern.""" + if len(w.shape) != 4: + return ops.ones_like(w) + _, _, (c_out, c_in, kh, kw) = kernels_and_patterns(w, src, epsilon=0.0) + _, distances = pattern_distances(w, dominant_patterns, valid_mask, src, epsilon, distance_metric) + closest = ops.argmin(distances, axis=1) # (M,) + mask_flat = ops.take(dominant_patterns, closest, axis=0) # (M, K) + mask_oihw = ops.reshape(mask_flat, (c_out, c_in, kh, kw)) + mask_src = convert_conv_layout(mask_oihw, src=CANONICAL_CONV_LAYOUT, dst=src) # back to weight layout + return ops.cast(mask_src, w.dtype) diff --git a/src/pquant/core/torch/pruning_methods/mdmm.py b/src/pquant/core/torch/pruning_methods/mdmm.py index 8afe615..3acfa53 100644 --- a/src/pquant/core/torch/pruning_methods/mdmm.py +++ b/src/pquant/core/torch/pruning_methods/mdmm.py @@ -9,6 +9,8 @@ LessThanOrEqualConstraint, ) from pquant.core.torch.pruning_methods.metric_functions import ( + FPGAAwareSparsityMetric, + PACAPatternMetric, StructuredSparsityMetric, UnstructuredSparsityMetric, ) @@ -16,6 +18,8 @@ _METRIC_REGISTRY = { "UnstructuredSparsity": UnstructuredSparsityMetric, "StructuredSparsity": StructuredSparsityMetric, + "FPGAAwareSparsity": FPGAAwareSparsityMetric, + "PACAPatternSparsity": PACAPatternMetric, } _CONSTRAINT_REGISTRY = { @@ -59,6 +63,14 @@ def build(self, input_shape): "l0_mode": l0_mode, "scale_mode": scale_mode, "rf": pruning_parameters.rf, + # FPGAAwareSparsityMetric + "precision": pruning_parameters.precision, + "target_resource": pruning_parameters.target_resource, + "bram_width": pruning_parameters.bram_width, + # PACAPatternMetric + "num_patterns_to_keep": pruning_parameters.num_patterns_to_keep, + "beta": pruning_parameters.beta, + "distance_metric": pruning_parameters.distance_metric, } metric_cls = _METRIC_REGISTRY.get(metric_type) @@ -85,9 +97,18 @@ def build(self, input_shape): self.register_buffer("mask", torch.ones(tuple(input_shape))) self.built = True + def _compute_hard_mask(self, weight, epsilon): + # During fine-tuning, a metric that defines its own projection (e.g. PACA pattern + # pruning) supplies the mask; otherwise use the magnitude threshold. The layer only + # checks for the capability, so it stays metric-agnostic (no metric-type branching). + metric_fn = getattr(self.constraint_layer, "metric_fn", None) + if self._is_finetuning and hasattr(metric_fn, "get_projection_mask"): + return metric_fn.get_projection_mask(weight).to(weight.dtype) + return (weight.abs() > epsilon).to(weight.dtype) + def forward(self, weight): epsilon = self.config.pruning_parameters.epsilon - hard_mask = (weight.abs() > epsilon).to(weight.dtype) + hard_mask = self._compute_hard_mask(weight, epsilon) not_active = self._is_pretraining or self._is_finetuning if not not_active: diff --git a/src/pquant/core/torch/pruning_methods/metric_functions.py b/src/pquant/core/torch/pruning_methods/metric_functions.py index e005c5c..1791cf7 100644 --- a/src/pquant/core/torch/pruning_methods/metric_functions.py +++ b/src/pquant/core/torch/pruning_methods/metric_functions.py @@ -1,5 +1,13 @@ import torch +from pquant.core.constants import ( + DISTANCE_VALUED_HAMMING, + FPGA_TARGET_RESOURCES, + TARGET_RESOURCE_BRAM, + TARGET_RESOURCE_DSP, +) +from pquant.core.torch.pruning_methods import patterns + class UnstructuredSparsityMetric: """L0-L1 based metric — torch port of the keras version.""" @@ -55,3 +63,115 @@ def __call__(self, weight): group_norms = torch.sqrt((groups.square()).sum(dim=-1)) zero_groups = (group_norms <= self.epsilon).to(torch.float32) return zero_groups.sum() / float(group_norms.numel()) + + +class FPGAAwareSparsityMetric: + """Hardware-aware sparsity metric for FPGA targets (torch port). + + Same semantics as the keras version: groups weights into DSP blocks (size ``rf``) and + BRAM blocks (``c`` DSP blocks, ``c`` derived from ``bram_width``/``precision``), and + returns the fraction of zero-valued groups at the chosen ``target_resource``. Inputs + are validated by the Pydantic config model, so no constructor asserts are kept here. + """ + + def __init__(self, rf=1, precision=16, target_resource=TARGET_RESOURCE_DSP, bram_width=36, epsilon=1e-3): + self.rf = int(rf) + self.precision = int(precision) + self.target_resource = target_resource + self.bram_width = int(bram_width) + self.epsilon = float(epsilon) + self.c = self._calculate_c() + + def _calculate_c(self): + """Number of consecutive DSP groups packed into a single BRAM block.""" + if self.bram_width % self.precision == 0: + return self.bram_width // self.precision + return (2 * self.bram_width) // self.precision + + def _prepare_weights(self, weight): + """Flatten to (rows, -1) and right-pad the row length to a multiple of rf.""" + if weight.dim() == 1: + w = weight.reshape(1, -1) + elif weight.dim() > 2: + w = weight.reshape(weight.shape[0], -1) + else: + w = weight + padding = (self.rf - w.shape[1] % self.rf) % self.rf + if padding: + w = torch.nn.functional.pad(w, (0, padding)) + return w + + def __call__(self, weight): + prepared = self._prepare_weights(weight) + dsp_groups = prepared.reshape(prepared.shape[0], -1, self.rf) + if self.target_resource == TARGET_RESOURCE_DSP: + return self._dsp_sparsity(dsp_groups) + if self.target_resource == TARGET_RESOURCE_BRAM: + return self._bram_sparsity(dsp_groups) + raise ValueError(f"target_resource must be one of {FPGA_TARGET_RESOURCES}, got {self.target_resource!r}") + + def _dsp_sparsity(self, dsp_groups): + group_norms = torch.sqrt(dsp_groups.square().sum(dim=-1)) + zero_groups = (group_norms <= self.epsilon).to(dsp_groups.dtype) + return zero_groups.sum() / float(group_norms.numel()) + + def _bram_sparsity(self, dsp_groups): + if self.c < 1: + raise ValueError( + f"BRAM packing needs precision <= 2*bram_width (got precision={self.precision}, " + f"bram_width={self.bram_width} -> c={self.c})." + ) + num_dsp_groups = dsp_groups.shape[1] + padding = (self.c - num_dsp_groups % self.c) % self.c + if padding: + dsp_groups = torch.nn.functional.pad(dsp_groups, (0, 0, 0, padding)) + bram_groups = dsp_groups.reshape(dsp_groups.shape[0], -1, self.c, self.rf) + bram_norms = torch.sqrt(bram_groups.square().sum(dim=(-1, -2))) + zero_bram = (bram_norms <= self.epsilon).to(bram_groups.dtype) + return zero_bram.sum() / float(bram_norms.numel()) + + +class PACAPatternMetric: + """Pattern-based pruning metric (PACA, torch port). + + On the first call it selects a small set of dominant binary patterns over the conv + kernels (cached for the metric instance lifetime), then returns the mean distance of + every kernel to its closest dominant pattern. Operates on 4D conv weights only (torch + kernels are OIHW); returns 0 for non-4D inputs. Exposes ``get_projection_mask`` so the + MDMM layer can snap weights onto their patterns during fine-tuning. + """ + + def __init__( + self, num_patterns_to_keep=16, beta=0.75, epsilon=1e-5, distance_metric=DISTANCE_VALUED_HAMMING, src="OIHW" + ): + self.num_patterns_to_keep = int(num_patterns_to_keep) + self.beta = float(beta) + self.epsilon = float(epsilon) + self.distance_metric = distance_metric + self.src = src # torch conv kernels are OIHW + self.dominant_patterns = None + self.valid_mask = None + + def _ensure_patterns(self, weight): + if self.dominant_patterns is None: + _, pats, _ = patterns.kernels_and_patterns(weight, self.src, self.epsilon) + self.dominant_patterns, self.valid_mask = patterns.select_dominant_patterns( + pats, self.num_patterns_to_keep, self.beta + ) + + def __call__(self, weight): + if weight.dim() != 4: + return torch.zeros((), dtype=weight.dtype, device=weight.device) + self._ensure_patterns(weight) + _, distances = patterns.pattern_distances( + weight, self.dominant_patterns, self.valid_mask, self.src, self.epsilon, self.distance_metric + ) + return distances.min(dim=1).values.mean() + + def get_projection_mask(self, weight): + # Identity mask if patterns were never selected so MDMM's `weight * mask` is a no-op. + if self.dominant_patterns is None: + return torch.ones_like(weight) + return patterns.projection_mask( + weight, self.dominant_patterns, self.valid_mask, self.src, self.epsilon, self.distance_metric + ) diff --git a/src/pquant/core/torch/pruning_methods/patterns.py b/src/pquant/core/torch/pruning_methods/patterns.py new file mode 100644 index 0000000..8343a26 --- /dev/null +++ b/src/pquant/core/torch/pruning_methods/patterns.py @@ -0,0 +1,126 @@ +# @Author: Arghya Ranjan Das +# PACA pattern utilities (Torch backend) — native-torch port of the keras version. +# +# Torch conv weights are already OIHW (the canonical layout), so layout conversion is +# normally a no-op; `src` is kept for API parity and correctness should a different layout +# ever be passed. Unique-pattern counting uses torch.unique (on-device, no NumPy). The +# dominant set is returned at a fixed size (num_patterns_to_keep) with a validity mask, so +# downstream shapes are static and the weight never leaves the device. + +import torch + +from pquant.core.constants import ( + CANONICAL_CONV_LAYOUT, + CONV_LAYOUT_AXES, + DISTANCE_COSINE, + DISTANCE_HAMMING, + DISTANCE_VALUED_HAMMING, +) + +_INF = 1e30 + + +def _layout_to_axes(layout): + if len(layout) != 4 or set(layout) != set("HWIO"): + raise ValueError(f"layout must be a permutation of 'HWIO', got {layout!r}") + return tuple(CONV_LAYOUT_AXES[ch] for ch in layout) + + +def _perm(src, dst): + """Permutation tuple that reorders axes from `src` layout to `dst` layout.""" + s = _layout_to_axes(src) + d = _layout_to_axes(dst) + return tuple(s.index(ax) for ax in d) + + +def convert_conv_layout(w, src, dst=CANONICAL_CONV_LAYOUT): + """Permute a 4D conv weight from `src` to `dst` layout (no-op if already equal).""" + if src == dst: + return w + perm = _perm(src, dst) + if perm == (0, 1, 2, 3): + return w + return w.permute(*perm).contiguous() + + +def kernels_and_patterns(w, src, epsilon): + """Flatten a 4D conv weight to per-kernel rows and their binary support. + + Returns (kernels, patterns, (C_out, C_in, kH, kW)): + kernels: (C_out*C_in, kH*kW) float - flattened kernels, canonical OIHW order. + patterns: (C_out*C_in, kH*kW) uint8 - binary support, |w| > epsilon. + """ + w_oihw = convert_conv_layout(w, src=src, dst=CANONICAL_CONV_LAYOUT) + c_out, c_in, kh, kw = w_oihw.shape + kernels = w_oihw.reshape(c_out * c_in, kh * kw) + patterns = (kernels.abs() > epsilon).to(torch.uint8) + return kernels, patterns, (c_out, c_in, kh, kw) + + +def select_dominant_patterns(patterns, num_patterns_to_keep, beta): + """Select the most frequent distinct patterns covering `beta` of the total mass. + + Native torch (torch.unique). Returns (dominant, valid): + dominant: (num_patterns_to_keep, kH*kW) uint8 - the candidate patterns. + valid: (num_patterns_to_keep,) bool - which rows are real selections + (padding rows past the beta cut / distinct-pattern count are ignored). + """ + alpha = int(num_patterns_to_keep) + m, k = patterns.shape + if m == 0: + return patterns.new_zeros((alpha, k)), patterns.new_zeros((alpha,), dtype=torch.bool) + + uniq, counts = torch.unique(patterns, dim=0, return_counts=True) # (U, K), (U,) + pdf = counts.to(torch.float32) / float(m) + order = torch.argsort(pdf, descending=True) + uniq, pdf = uniq[order], pdf[order] + cdf = torch.cumsum(pdf, dim=0) + + reaches = cdf >= beta + n_beta = int(torch.argmax(reaches.to(torch.int32))) + 1 if bool(reaches.any()) else uniq.shape[0] + keep = min(n_beta, uniq.shape[0], alpha) + + dom = uniq[:alpha] + if dom.shape[0] < alpha: + dom = torch.cat([dom, dom.new_zeros((alpha - dom.shape[0], k))], dim=0) + valid = torch.arange(alpha, device=patterns.device) < keep + return dom, valid + + +def _kernel_pattern_distances(kernel_patterns, kernels, dominant_patterns, distance_metric): + """Distance from every kernel to every dominant pattern. -> (M_kernels, alpha).""" + tk = kernel_patterns.to(kernels.dtype).unsqueeze(1) # (M, 1, K) + k_e = kernels.unsqueeze(1) # (M, 1, K) + p_e = dominant_patterns.to(kernels.dtype).unsqueeze(0) # (1, alpha, K) + + if distance_metric == DISTANCE_HAMMING: + return (tk - p_e).abs().sum(dim=-1) + if distance_metric == DISTANCE_VALUED_HAMMING: + return ((tk - p_e).abs() * k_e.abs()).sum(dim=-1) + if distance_metric == DISTANCE_COSINE: + projected = k_e * p_e + dot = (k_e * projected).sum(dim=-1) + denom = k_e.norm(dim=-1) * projected.norm(dim=-1) + 1e-7 + return 1.0 - dot / denom + raise ValueError(f"Unsupported distance metric: {distance_metric!r}") + + +def pattern_distances(w, dominant_patterns, valid_mask, src, epsilon, distance_metric): + """Per-kernel distance to each dominant pattern, with invalid patterns masked to +inf.""" + kernels, kernel_patterns, _ = kernels_and_patterns(w, src, epsilon) + distances = _kernel_pattern_distances(kernel_patterns, kernels, dominant_patterns, distance_metric) + distances = torch.where(valid_mask.unsqueeze(0), distances, distances.new_full((), _INF)) + return kernels, distances + + +def projection_mask(w, dominant_patterns, valid_mask, src, epsilon, distance_metric): + """Binary mask (same layout as `w`) projecting each kernel onto its closest dominant pattern.""" + if w.dim() != 4: + return torch.ones_like(w) + _, _, (c_out, c_in, kh, kw) = kernels_and_patterns(w, src, epsilon=0.0) + _, distances = pattern_distances(w, dominant_patterns, valid_mask, src, epsilon, distance_metric) + closest = torch.argmin(distances, dim=1) # (M,) + mask_flat = dominant_patterns[closest] # (M, K) + mask_oihw = mask_flat.reshape(c_out, c_in, kh, kw) + mask_src = convert_conv_layout(mask_oihw, src=CANONICAL_CONV_LAYOUT, dst=src) + return mask_src.to(w.dtype) diff --git a/src/pquant/data_models/pruning_model.py b/src/pquant/data_models/pruning_model.py index 21e6607..3e1f6ae 100644 --- a/src/pquant/data_models/pruning_model.py +++ b/src/pquant/data_models/pruning_model.py @@ -1,7 +1,7 @@ from enum import Enum from typing import List, Literal, Optional -from pydantic import BaseModel, Field +from pydantic import BaseModel, Field, model_validator class BasePruningModel(BaseModel): @@ -67,6 +67,8 @@ class ActivationPruningModel(BasePruningModel): class MetricType(str, Enum): UNSTRUCTURED = "UnstructuredSparsity" STRUCTURED = "StructuredSparsity" + FPGA_AWARE = "FPGAAwareSparsity" + PACA_PATTERN = "PACAPatternSparsity" class ConstraintType(str, Enum): @@ -77,9 +79,14 @@ class ConstraintType(str, Enum): class MDMMPruningModel(BasePruningModel): pruning_method: Literal["mdmm"] = "mdmm" - constraint_type: ConstraintType = Field("Equality") + # Defaults must be the enum members, not bare strings: Pydantic does not validate + # defaults, so a str default leaves the field holding a str at runtime and triggers + # `PydanticSerializationUnexpectedValue` ("Expected enum") on model_dump — which in + # turn makes the discriminated pruning_parameters union fall back to per-member + # serialization warnings. Using the enum members keeps serialization clean. + constraint_type: ConstraintType = Field(default=ConstraintType.EQUALITY) target_value: float = Field(default=0.0) - metric_type: MetricType = Field(default="UnstructuredSparsity") + metric_type: MetricType = Field(default=MetricType.UNSTRUCTURED) target_sparsity: float = Field(default=0.9) rf: int = Field(default=1) epsilon: float = Field(default=1.0e-03) @@ -89,3 +96,26 @@ class MDMMPruningModel(BasePruningModel): l0_mode: Literal["coarse", "smooth"] = Field(default="coarse") scale_mode: Literal["mean", "sum"] = Field(default="mean") constraint_lr: float = Field(default=1.0e-3) + # --- Hardware-aware metric parameters --- + # Each group is consumed only when the matching metric_type is selected. The MDMM + # layer filters these against the chosen metric's __init__ signature, so leaving a + # value as None falls back to that metric's own default. Validation lives here + # (ge/le + Literal) rather than as asserts inside the metric classes. + # FPGAAwareSparsityMetric (metric_type == "FPGAAwareSparsity") + precision: Optional[int] = Field(default=None, ge=1) + target_resource: Optional[Literal["DSP", "BRAM"]] = Field(default=None) + bram_width: Optional[int] = Field(default=None, ge=1) + # PACAPatternMetric (metric_type == "PACAPatternSparsity") + num_patterns_to_keep: Optional[int] = Field(default=None, ge=1) + beta: Optional[float] = Field(default=None, ge=0.0, le=1.0) + distance_metric: Optional[Literal["hamming", "valued_hamming", "cosine"]] = Field(default=None) + + @model_validator(mode="after") + def _enforce_paca_constraint(self): + # PACA pattern pruning is defined as driving the pattern-distance metric to + # zero, so it always pairs with an equality constraint at target 0. Enforced + # here (in the data model) so the MDMM layer can stay metric-agnostic. + if self.metric_type == MetricType.PACA_PATTERN: + self.constraint_type = ConstraintType.EQUALITY + self.target_value = 0.0 + return self diff --git a/tests/run_tests.sh b/tests/run_tests.sh index 66119ec..d3548e2 100755 --- a/tests/run_tests.sh +++ b/tests/run_tests.sh @@ -4,6 +4,9 @@ pytest test_ap.py KERAS_BACKEND="torch" pytest test_ap.py pytest test_pdp.py KERAS_BACKEND="torch" pytest test_pdp.py +pytest test_mdmm.py test_mdmm_metrics.py +KERAS_BACKEND="torch" pytest test_mdmm.py test_mdmm_metrics.py +KERAS_BACKEND="torch" pytest test_torch_mdmm.py pytest test_wanda.py KERAS_BACKEND="torch" pytest test_wanda.py pytest test_keras_compression_layers.py diff --git a/tests/test_mdmm.py b/tests/test_mdmm.py new file mode 100644 index 0000000..a150cc8 --- /dev/null +++ b/tests/test_mdmm.py @@ -0,0 +1,345 @@ +"""MDMM core behaviour (Keras layer): metrics, constraints, masking and training phases. + +Imports the Keras MDMM implementation but uses backend-agnostic keras.ops, so the same +file runs under both KERAS_BACKEND=tensorflow and KERAS_BACKEND=torch (see run_tests.sh). +Native-torch MDMM is covered separately in test_torch_mdmm.py. +""" + +import math + +import numpy as np +import pytest +from keras import ops + +from pquant.core.keras.pruning_methods.constraint_functions import ( + EqualityConstraint, + GreaterThanOrEqualConstraint, + LessThanOrEqualConstraint, +) +from pquant.core.keras.pruning_methods.mdmm import MDMM +from pquant.core.keras.pruning_methods.metric_functions import UnstructuredSparsityMetric + +IN_FEATURES = 8 +OUT_FEATURES = 16 + + +@pytest.fixture +def config(): + return { + "pruning_parameters": { + "pruning_method": "mdmm", + "disable_pruning_for_layers": [], + "enable_pruning": True, + "constraint_type": "Equality", + "target_value": 0.0, + "metric_type": "UnstructuredSparsity", + "target_sparsity": 0.5, + "rf": 1, + "epsilon": 1e-3, + "scale": 1.0, + "damping": 1.0, + "use_grad": False, + "l0_mode": "coarse", + "scale_mode": "mean", + "constraint_lr": 0.0, + } + } + + +# --- UnstructuredSparsityMetric in isolation --- + + +def test_unstructured_sparsity_coarse_mean_satisfied(): + """weight=[0,0,1,2], target=0.5 -> L0=0.5, factor=0, metric=0.""" + metric = UnstructuredSparsityMetric(l0_mode="coarse", scale_mode="mean", epsilon=1e-3, target_sparsity=0.5) + weight = ops.convert_to_tensor([0.0, 0.0, 1.0, 2.0], dtype="float32") + assert ops.abs(metric(weight)) < 1e-6 + + +def test_unstructured_sparsity_coarse_mean_unsatisfied(): + """weight=[0,0,0,2], target=0.5 -> L0=0.75, factor=-0.3125, metric=-0.15625.""" + metric = UnstructuredSparsityMetric(l0_mode="coarse", scale_mode="mean", epsilon=1e-3, target_sparsity=0.5) + weight = ops.convert_to_tensor([0.0, 0.0, 0.0, 2.0], dtype="float32") + assert ops.abs(metric(weight) - (-0.15625)) < 1e-5 + + +def test_unstructured_sparsity_coarse_sum(): + """Same as unsatisfied but scale_mode='sum' -> no division by num_weights.""" + metric = UnstructuredSparsityMetric(l0_mode="coarse", scale_mode="sum", epsilon=1e-3, target_sparsity=0.5) + weight = ops.convert_to_tensor([0.0, 0.0, 0.0, 2.0], dtype="float32") + assert ops.abs(metric(weight) - (-0.625)) < 1e-5 + + +def test_unstructured_sparsity_smooth_at_zeros(): + """weight=[0,0,1,2], target=0.5 -> smooth L0 ~ 0.5, metric ~ 0.""" + metric = UnstructuredSparsityMetric(l0_mode="smooth", scale_mode="mean", epsilon=1e-3, target_sparsity=0.5) + weight = ops.convert_to_tensor([0.0, 0.0, 1.0, 2.0], dtype="float32") + assert ops.abs(metric(weight)) < 1e-4 + + +def test_unstructured_sparsity_smooth_nonzero(): + """weight=[0.1,0.1,1,2], target=0.5 -> smooth L0 ~ 0.184, metric ~ 0.173.""" + metric = UnstructuredSparsityMetric(l0_mode="smooth", scale_mode="mean", epsilon=1e-3, target_sparsity=0.5) + weight = ops.convert_to_tensor([0.1, 0.1, 1.0, 2.0], dtype="float32") + smooth_l0 = (2 * math.exp(-1.0)) / 4.0 + expected = (0.5**2 - smooth_l0**2) * 3.2 / 4.0 + assert ops.abs(metric(weight) - expected) < 1e-4 + + +# --- Constraint functions in isolation --- + + +class ConstantMetric: + """A metric that always returns a fixed value, to isolate constraint logic.""" + + def __init__(self, value): + self.value = value + + def __call__(self, weight): + return ops.convert_to_tensor(self.value, dtype="float32") + + +def test_equality_constraint(): + """metric=0.3, target=0.0 -> inf=0.3, penalty=1*(0.3 + 0.09/2)=0.345.""" + constraint = EqualityConstraint(metric_fn=ConstantMetric(0.3), target_value=0.0, scale=1.0, damping=1.0, + use_grad=False, lr=0.0) + dummy = ops.zeros((2, 2)) + constraint.build(dummy.shape) + assert ops.abs(constraint(dummy) - 0.345) < 1e-5 + + +def test_leq_constraint_satisfied(): + constraint = LessThanOrEqualConstraint(metric_fn=ConstantMetric(0.3), target_value=0.5, scale=1.0, damping=1.0, + use_grad=False, lr=0.0) + dummy = ops.zeros((2, 2)) + constraint.build(dummy.shape) + assert ops.abs(constraint(dummy)) < 1e-6 + + +def test_leq_constraint_violated(): + """metric=0.7, target=0.5 -> inf=0.2, penalty=0.22.""" + constraint = LessThanOrEqualConstraint(metric_fn=ConstantMetric(0.7), target_value=0.5, scale=1.0, damping=1.0, + use_grad=False, lr=0.0) + dummy = ops.zeros((2, 2)) + constraint.build(dummy.shape) + assert ops.abs(constraint(dummy) - 0.22) < 1e-5 + + +def test_geq_constraint_satisfied(): + constraint = GreaterThanOrEqualConstraint(metric_fn=ConstantMetric(0.7), target_value=0.5, scale=1.0, damping=1.0, + use_grad=False, lr=0.0) + dummy = ops.zeros((2, 2)) + constraint.build(dummy.shape) + assert ops.abs(constraint(dummy)) < 1e-6 + + +def test_geq_constraint_violated(): + """metric=0.3, target=0.5 -> inf=0.2, penalty=0.22.""" + constraint = GreaterThanOrEqualConstraint(metric_fn=ConstantMetric(0.3), target_value=0.5, scale=1.0, damping=1.0, + use_grad=False, lr=0.0) + dummy = ops.zeros((2, 2)) + constraint.build(dummy.shape) + assert ops.abs(constraint(dummy) - 0.22) < 1e-5 + + +def test_constraint_turn_off(): + """After turn_off(), penalty should be 0.""" + constraint = EqualityConstraint(metric_fn=ConstantMetric(0.5), target_value=0.0, scale=10.0, damping=5.0, + use_grad=False, lr=0.0) + dummy = ops.zeros((2, 2)) + constraint.build(dummy.shape) + constraint.turn_off() + assert ops.abs(constraint(dummy)) < 1e-6 + + +def test_constraint_lambda_update_no_grad(): + """With lr>0, lambda updates between calls, changing the penalty.""" + constraint = EqualityConstraint(metric_fn=ConstantMetric(0.3), target_value=0.0, scale=2.0, damping=1.0, + use_grad=False, lr=0.1) + dummy = ops.zeros((2, 2)) + constraint.build(dummy.shape) + result1 = constraint(dummy, training=True) + result2 = constraint(dummy, training=True) + assert not ops.isclose(result1, result2) + + +# --- Mask correctness --- + + +def test_hard_mask_threshold(config): + """|w| > epsilon -> 1, else 0 (mask updated during active phase).""" + weight = ops.convert_to_tensor([[1e-4, 1e-3, 1e-2, 0.1], [-1e-4, -1e-3, -1e-2, -0.1]], dtype="float32") + expected_mask = ops.convert_to_tensor([[0.0, 0.0, 1.0, 1.0], [0.0, 0.0, 1.0, 1.0]], dtype="float32") + mdmm = MDMM(config, "linear") + mdmm.build(weight.shape) + mdmm.post_pre_train_function() + mdmm(weight) + assert ops.all(ops.equal(ops.convert_to_tensor(mdmm.mask), expected_mask)) + + +def test_get_layer_sparsity(config): + """Half above epsilon -> keep ratio 0.5 (also guards the float/int dtype fix).""" + weight = ops.convert_to_tensor([[1e-4, 1e-3, 1e-2, 0.1], [-1e-4, -1e-3, -1e-2, -0.1]], dtype="float32") + mdmm = MDMM(config, "linear") + mdmm.build(weight.shape) + assert ops.abs(mdmm.get_layer_sparsity(weight) - 0.5) < 1e-5 + + +# --- MDMM training phases --- + + +def _make_mixed_weight_linear(): + """~75% zeros (below epsilon), ~25% nonzero, so L0 != target_sparsity=0.5.""" + vals = np.zeros(OUT_FEATURES * IN_FEATURES, dtype=np.float32) + quarter = len(vals) // 4 + vals[-quarter:] = np.linspace(0.01, 1.0, quarter) + return ops.convert_to_tensor(vals.reshape(OUT_FEATURES, IN_FEATURES)) + + +def test_pretraining_phase_linear(config): + """Pretraining: mask stays ones, no penalty, output unchanged.""" + weight = _make_mixed_weight_linear() + mdmm = MDMM(config, "linear") + mdmm.build(weight.shape) + result = mdmm(weight) + assert ops.all(ops.equal(result, weight)) + assert ops.all(ops.equal(ops.convert_to_tensor(mdmm.mask), ops.ones(weight.shape))) + assert mdmm.losses[-1] < 1e-8 + + +def test_active_phase_linear(config): + """Active: mask updated, penalty nonzero, output unchanged.""" + weight = _make_mixed_weight_linear() + mdmm = MDMM(config, "linear") + mdmm.build(weight.shape) + mdmm.post_pre_train_function() + result = mdmm(weight) + expected_mask = ops.cast(ops.abs(weight) > config["pruning_parameters"]["epsilon"], weight.dtype) + assert ops.all(ops.equal(result, weight)) + assert ops.all(ops.equal(ops.convert_to_tensor(mdmm.mask), expected_mask)) + assert mdmm.losses[-1] > 1e-8 + + +def test_finetuning_phase_linear(config): + """Finetuning: mask frozen, no penalty, output = weight * hard_mask.""" + weight = _make_mixed_weight_linear() + epsilon = config["pruning_parameters"]["epsilon"] + mdmm = MDMM(config, "linear") + mdmm.build(weight.shape) + mdmm.post_pre_train_function() + mdmm(weight) + mdmm.pre_finetune_function() + result = mdmm(weight) + expected_output = weight * ops.cast(ops.abs(weight) > epsilon, weight.dtype) + assert ops.all(ops.equal(result, expected_output)) + assert mdmm.losses[-1] < 1e-8 + + +# --- Penalty numerical verification (full chain) --- + + +def test_penalty_satisfied(config): + """weight=[0,0,1,2] at target=0.5 -> metric=0 -> penalty=0.""" + config["pruning_parameters"]["target_sparsity"] = 0.5 + weight = ops.convert_to_tensor([[0.0, 0.0], [1.0, 2.0]], dtype="float32") + mdmm = MDMM(config, "linear") + mdmm.build(weight.shape) + mdmm.post_pre_train_function() + mdmm(weight) + assert mdmm.losses[-1] < 1e-6 + + +def test_penalty_unsatisfied(config): + """weight=[0,0,0,2] at target=0.5 -> known penalty value.""" + config["pruning_parameters"]["target_sparsity"] = 0.5 + weight = ops.convert_to_tensor([[0.0, 0.0], [0.0, 2.0]], dtype="float32") + mdmm = MDMM(config, "linear") + mdmm.build(weight.shape) + mdmm.post_pre_train_function() + mdmm(weight) + expected = 0.15625 + 0.15625**2 / 2.0 + assert ops.abs(mdmm.losses[-1] - expected) < 1e-4 + + +def test_penalty_with_scale_and_damping(config): + """Same weight but scale=10, damping=2 -> larger penalty.""" + config["pruning_parameters"]["target_sparsity"] = 0.5 + config["pruning_parameters"]["scale"] = 10.0 + config["pruning_parameters"]["damping"] = 2.0 + weight = ops.convert_to_tensor([[0.0, 0.0], [0.0, 2.0]], dtype="float32") + mdmm = MDMM(config, "linear") + mdmm.build(weight.shape) + mdmm.post_pre_train_function() + mdmm(weight) + inf_val = 0.15625 + expected = 10.0 * (1.0 * inf_val + 2.0 * inf_val**2 / 2.0) + assert ops.abs(mdmm.losses[-1] - expected) < 1e-3 + + +# --- Constraint-type variations (full chain) --- + + +def test_leq_full_chain_satisfied(config): + """LEQ with metric < 0 (satisfied) -> penalty ~ 0.""" + config["pruning_parameters"]["constraint_type"] = "LessThanOrEqual" + config["pruning_parameters"]["target_sparsity"] = 0.5 + weight = ops.convert_to_tensor([[0.0, 0.0], [0.0, 2.0]], dtype="float32") + mdmm = MDMM(config, "linear") + mdmm.build(weight.shape) + mdmm.post_pre_train_function() + mdmm(weight) + assert mdmm.losses[-1] < 1e-6 + + +def test_geq_full_chain_violated(config): + """GEQ with metric < 0 (violated) -> nonzero penalty.""" + config["pruning_parameters"]["constraint_type"] = "GreaterThanOrEqual" + config["pruning_parameters"]["target_sparsity"] = 0.5 + weight = ops.convert_to_tensor([[0.0, 0.0], [0.0, 2.0]], dtype="float32") + mdmm = MDMM(config, "linear") + mdmm.build(weight.shape) + mdmm.post_pre_train_function() + mdmm(weight) + assert mdmm.losses[-1] > 1e-4 + + +def test_smooth_differs_from_coarse_in_soft_region(config): + """For w=0.05, smooth L0 differs from coarse -> different penalties.""" + weight = ops.convert_to_tensor([[0.05, 0.05], [1.0, 2.0]], dtype="float32") + + config["pruning_parameters"]["l0_mode"] = "coarse" + mdmm_coarse = MDMM(config, "linear") + mdmm_coarse.build(weight.shape) + mdmm_coarse.post_pre_train_function() + mdmm_coarse(weight) + + config["pruning_parameters"]["l0_mode"] = "smooth" + mdmm_smooth = MDMM(config, "linear") + mdmm_smooth.build(weight.shape) + mdmm_smooth.post_pre_train_function() + mdmm_smooth(weight) + + assert not ops.isclose(mdmm_coarse.losses[-1], mdmm_smooth.losses[-1]) + + +# --- Edge cases --- + + +def test_all_zero_weights(config): + """All-zero weights -> L1=0 -> metric=0 -> penalty=0.""" + weight = ops.zeros((4, 4)) + mdmm = MDMM(config, "linear") + mdmm.build(weight.shape) + mdmm.post_pre_train_function() + mdmm(weight) + assert mdmm.losses[-1] < 1e-6 + + +def test_all_large_weights(config): + """No weights below epsilon -> L0=0 -> factor=target^2 > 0 -> nonzero penalty.""" + weight = ops.ones((4, 4)) + mdmm = MDMM(config, "linear") + mdmm.build(weight.shape) + mdmm.post_pre_train_function() + mdmm(weight) + assert mdmm.losses[-1] > 1e-4 diff --git a/tests/test_mdmm_metrics.py b/tests/test_mdmm_metrics.py new file mode 100644 index 0000000..984a648 --- /dev/null +++ b/tests/test_mdmm_metrics.py @@ -0,0 +1,289 @@ +"""Hardware-aware MDMM metrics (Keras): FPGAAwareSparsity, PACAPattern + pattern helpers. + +Runs under both KERAS_BACKEND=tensorflow and torch (see run_tests.sh). The native-torch +metric implementations are exercised in test_torch_mdmm.py. Conv weights here use the Keras +HWIO kernel layout. +""" + +import math + +import numpy as np +import pytest +from keras import ops +from pydantic import ValidationError + +from pquant.core.keras.pruning_methods import patterns +from pquant.core.keras.pruning_methods.constraint_functions import EqualityConstraint +from pquant.core.keras.pruning_methods.mdmm import MDMM +from pquant.core.keras.pruning_methods.metric_functions import FPGAAwareSparsityMetric, PACAPatternMetric +from pquant.data_models.pruning_model import MDMMPruningModel + + +@pytest.fixture +def base_config(): + return { + "pruning_parameters": { + "pruning_method": "mdmm", + "disable_pruning_for_layers": [], + "enable_pruning": True, + "constraint_type": "Equality", + "target_value": 0.0, + "metric_type": "UnstructuredSparsity", + "target_sparsity": 0.5, + "rf": 1, + "epsilon": 1e-3, + "scale": 1.0, + "damping": 1.0, + "use_grad": False, + "constraint_lr": 0.0, + } + } + + +def _make_conv_weight(kh=3, kw=3, c_in=2, c_out=4, seed=42): + """Keras conv kernel: HWIO layout (kH, kW, C_in, C_out).""" + rng = np.random.default_rng(seed) + return ops.convert_to_tensor(rng.standard_normal((kh, kw, c_in, c_out)).astype(np.float32)) + + +def _conv_weight_for_mdmm(): + """HWIO weight with half the (out-channel) kernels zeroed for a non-trivial mask.""" + rng = np.random.default_rng(0) + arr = rng.standard_normal((3, 3, 4, 4)).astype(np.float32) + arr[..., :2] = 0.0 # zero half the output channels + return ops.convert_to_tensor(arr) + + +# --- FPGAAwareSparsityMetric --- + + +@pytest.mark.parametrize( + "target_resource,fill,expected", + [("DSP", 0.0, 1.0), ("DSP", 1.0, 0.0), ("BRAM", 0.0, 1.0), ("BRAM", 1.0, 0.0)], +) +def test_fpga_extreme_weights(target_resource, fill, expected): + metric = FPGAAwareSparsityMetric(rf=2, precision=16, bram_width=36, target_resource=target_resource, epsilon=1e-3) + weight = ops.cast(ops.full((4, 16), fill), "float32") + assert ops.abs(metric(weight) - expected) < 1e-5 + + +def test_fpga_dsp_l2_grouping_math(): + """[[1,1,0,0]] with rf=2 -> groups [1,1] (norm sqrt 2) and [0,0] (norm 0). Result = 0.5.""" + metric = FPGAAwareSparsityMetric(rf=2, target_resource="DSP", epsilon=1e-3) + weight = ops.convert_to_tensor([[1.0, 1.0, 0.0, 0.0]], dtype="float32") + assert ops.abs(metric(weight) - 0.5) < 1e-5 + + +@pytest.mark.parametrize("precision,bram_width,expected_c", [(16, 36, 4), (18, 36, 2)]) +def test_fpga_calculate_c(precision, bram_width, expected_c): + metric = FPGAAwareSparsityMetric(precision=precision, bram_width=bram_width, target_resource="DSP") + assert metric.c == expected_c + + +def test_fpga_handles_1d_weight(): + """1D bias vector reshapes to (1, N) and works without crashing.""" + metric = FPGAAwareSparsityMetric(rf=2, target_resource="DSP", epsilon=1e-3) + assert ops.abs(metric(ops.ones((8,), dtype="float32"))) < 1e-5 + + +def test_fpga_dtype_propagates_for_float64(): + """float64 weight returns a float64 scalar (not hardcoded float32).""" + metric = FPGAAwareSparsityMetric(rf=2, target_resource="DSP", epsilon=1e-3) + result = metric(ops.cast(ops.zeros((4, 8)), "float64")) + assert "float64" in str(result.dtype) + + +def test_fpga_invalid_target_resource_raises_on_call(): + """Validation now lives in Pydantic, so a bad target_resource surfaces at call time.""" + metric = FPGAAwareSparsityMetric(target_resource="LUT") + with pytest.raises(ValueError): + metric(ops.ones((4, 8), dtype="float32")) + + +def test_fpga_bram_packing_impossible_raises_on_call(): + """precision > 2*bram_width -> c=0 -> clear error when BRAM sparsity is computed.""" + metric = FPGAAwareSparsityMetric(precision=128, bram_width=16, target_resource="BRAM") + with pytest.raises(ValueError): + metric(ops.ones((4, 16), dtype="float32")) + + +# --- PACAPatternMetric --- + + +def test_paca_returns_zero_for_non_4d(): + metric = PACAPatternMetric() + weight = ops.convert_to_tensor(np.random.rand(8, 4).astype(np.float32)) + assert ops.abs(metric(weight)) < 1e-6 + + +def test_paca_dominant_patterns_lazy_caching(): + metric = PACAPatternMetric(num_patterns_to_keep=4, beta=0.75) + weight = _make_conv_weight() + assert metric.dominant_patterns is None + _ = metric(weight) + first = metric.dominant_patterns + assert first is not None + _ = metric(weight) + assert metric.dominant_patterns is first # cached, not recomputed + + +def test_paca_get_projection_mask_shape_and_binary(): + metric = PACAPatternMetric(num_patterns_to_keep=4, beta=0.75) + weight = _make_conv_weight() + _ = metric(weight) + mask = metric.get_projection_mask(weight) + assert tuple(ops.shape(mask)) == tuple(ops.shape(weight)) + assert set(np.unique(ops.convert_to_numpy(mask)).tolist()).issubset({0, 1}) + + +@pytest.mark.parametrize("distance_metric", ["hamming", "valued_hamming", "cosine"]) +def test_paca_all_distance_metrics(distance_metric): + metric = PACAPatternMetric(num_patterns_to_keep=4, beta=0.75, distance_metric=distance_metric) + assert math.isfinite(float(ops.convert_to_numpy(metric(_make_conv_weight())))) + + +def test_paca_projection_mask_identity_when_no_patterns(): + metric = PACAPatternMetric(num_patterns_to_keep=4, beta=0.75) + weight = _make_conv_weight() + mask = metric.get_projection_mask(weight) # never called -> identity + np.testing.assert_array_equal(ops.convert_to_numpy(mask), np.ones_like(ops.convert_to_numpy(mask))) + + +def test_paca_invalid_distance_metric_raises_on_call(): + metric = PACAPatternMetric(distance_metric="manhattan") + with pytest.raises(ValueError): + metric(_make_conv_weight()) + + +# --- Pattern helpers --- + + +def test_patterns_kernels_and_patterns(): + """Flatten per-kernel and binarize by epsilon (explicit OIHW input).""" + weight = ops.convert_to_tensor(np.array([[[[1.0, 0.0], [0.0, 1.0]]]], dtype=np.float32)) # (1,1,2,2) + kernels, all_patterns, _ = patterns.kernels_and_patterns(weight, src="OIHW", epsilon=0.5) + assert tuple(ops.shape(kernels)) == (1, 4) + np.testing.assert_array_equal(ops.convert_to_numpy(all_patterns), np.array([[1, 0, 0, 1]], dtype=np.uint8)) + + +def test_patterns_select_dominant_returns_fixed_size_with_validity(): + """select_dominant_patterns returns (num_patterns_to_keep, K) + a validity mask.""" + # 6 identical-support rows + 2 of another support -> 2 distinct patterns. + pats = ops.convert_to_tensor( + np.array([[1, 1, 0, 0]] * 6 + [[0, 0, 1, 1]] * 2, dtype=np.uint8) + ) + dom, valid = patterns.select_dominant_patterns(pats, num_patterns_to_keep=4, beta=0.99) + assert tuple(ops.shape(dom)) == (4, 4) + valid_np = ops.convert_to_numpy(valid) + assert valid_np[:2].all() and not valid_np[2:].any() # exactly 2 distinct patterns are valid + + +# --- MDMM integration through the registry --- + + +def test_mdmm_fpga_full_phase_cycle(base_config): + base_config["pruning_parameters"].update( + {"metric_type": "FPGAAwareSparsity", "precision": 16, "target_resource": "DSP", "bram_width": 36, "rf": 2} + ) + weight = _conv_weight_for_mdmm() + mdmm = MDMM(base_config, "conv") + mdmm.build(weight.shape) + + out = mdmm(weight) # pretraining + assert ops.all(ops.equal(out, weight)) + assert mdmm.losses[-1] < 1e-8 + + mdmm.post_pre_train_function() + _ = mdmm(weight) # active + mdmm.pre_finetune_function() + _ = mdmm(weight) # finetuning + assert mdmm.losses[-1] < 1e-8 + + +def test_mdmm_paca_forces_equality_constraint(base_config): + """PACA always pairs with EqualityConstraint(target=0), even if config asks for GEQ/99.""" + base_config["pruning_parameters"].update( + {"metric_type": "PACAPatternSparsity", "constraint_type": "GreaterThanOrEqual", "target_value": 99.0, + "num_patterns_to_keep": 4, "beta": 0.85, "distance_metric": "cosine"} + ) + mdmm = MDMM(base_config, "conv") + mdmm.build(_conv_weight_for_mdmm().shape) + assert isinstance(mdmm.constraint_layer, EqualityConstraint) + assert mdmm.constraint_layer.target_value == 0.0 + + +def test_mdmm_paca_full_phase_cycle(base_config): + """PACA: patterns populated once the metric is called; finetuning output = weight * binary mask.""" + base_config["pruning_parameters"].update( + {"metric_type": "PACAPatternSparsity", "num_patterns_to_keep": 4, "beta": 0.85, "distance_metric": "cosine"} + ) + weight = _conv_weight_for_mdmm() + mdmm = MDMM(base_config, "conv") + mdmm.build(weight.shape) + + out = mdmm(weight) # pretraining (constraint still evaluates the metric -> selects patterns) + assert ops.all(ops.equal(out, weight)) + assert mdmm.constraint_layer.metric_fn.dominant_patterns is not None + + mdmm.post_pre_train_function() + _ = mdmm(weight) # active + mdmm.pre_finetune_function() + out = mdmm(weight) # finetuning + out_np, weight_np = ops.convert_to_numpy(out), ops.convert_to_numpy(weight) + assert np.logical_or(np.isclose(out_np, 0.0), np.isclose(out_np, weight_np)).all() + + +# --- Pydantic schema validation (backend-agnostic) --- + + +def test_pydantic_accepts_new_metric_types(): + fpga = MDMMPruningModel(metric_type="FPGAAwareSparsity", precision=16, target_resource="DSP", bram_width=36) + assert fpga.metric_type.value == "FPGAAwareSparsity" and fpga.precision == 16 + paca = MDMMPruningModel(metric_type="PACAPatternSparsity", num_patterns_to_keep=8, beta=0.9, distance_metric="hamming") + assert paca.metric_type.value == "PACAPatternSparsity" and paca.beta == 0.9 + + +@pytest.mark.parametrize( + "kwargs", + [ + {"metric_type": "BogusMetric"}, + {"target_resource": "LUT"}, + {"distance_metric": "manhattan"}, + {"precision": 0}, + {"bram_width": 0}, + {"num_patterns_to_keep": 0}, + {"beta": 1.5}, + {"beta": -0.1}, + ], +) +def test_pydantic_rejects_invalid_values(kwargs): + with pytest.raises(ValidationError): + MDMMPruningModel(**kwargs) + + +def test_pydantic_paca_forces_equality_at_model_level(): + """The model_validator coerces PACA to Equality/0 regardless of requested constraint.""" + m = MDMMPruningModel(metric_type="PACAPatternSparsity", constraint_type="GreaterThanOrEqual", target_value=42.0) + assert m.constraint_type.value == "Equality" and m.target_value == 0.0 + + +def test_mdmm_config_serializes_without_warnings(): + """Regression for the reported bug: MDMM configs must round-trip through model_dump + with no PydanticSerializationUnexpectedValue (constraint_type enum / discriminated union).""" + import warnings + + from pquant.core.hyperparameter_optimization import PQConfig + + cases = [ + {}, + {"metric_type": "FPGAAwareSparsity", "precision": 8, "target_resource": "DSP"}, + {"metric_type": "PACAPatternSparsity", "num_patterns_to_keep": 8}, + ] + for pp in cases: + cfg = PQConfig.load_from_config({"pruning_parameters": {"pruning_method": "mdmm", **pp}}) + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always") + cfg.model_dump(mode="json") + cfg.model_dump_json() + serial = [w for w in caught if "Serialization" in str(w.message) or "serialized value" in str(w.message)] + assert not serial, f"serialization warnings for {pp}: {[str(w.message) for w in serial]}" diff --git a/tests/test_torch_mdmm.py b/tests/test_torch_mdmm.py new file mode 100644 index 0000000..07c84f2 --- /dev/null +++ b/tests/test_torch_mdmm.py @@ -0,0 +1,207 @@ +"""Native-torch MDMM implementation: FPGA/PACA metrics, pattern helpers and phase cycle. + +Exercises pquant.core.torch.pruning_methods directly (raw torch.nn.Module, native torch +ops). Run with KERAS_BACKEND=torch (see run_tests.sh). Torch conv kernels are OIHW. +""" + +import math + +import numpy as np +import pytest +import torch + +from pquant.core.torch.pruning_methods import patterns +from pquant.core.torch.pruning_methods.constraint_functions import EqualityConstraint +from pquant.core.torch.pruning_methods.mdmm import MDMM +from pquant.core.torch.pruning_methods.metric_functions import ( + FPGAAwareSparsityMetric, + PACAPatternMetric, + UnstructuredSparsityMetric, +) + + +def base_config(): + return { + "pruning_parameters": { + "pruning_method": "mdmm", + "disable_pruning_for_layers": [], + "enable_pruning": True, + "constraint_type": "Equality", + "target_value": 0.0, + "metric_type": "UnstructuredSparsity", + "target_sparsity": 0.5, + "rf": 1, + "epsilon": 1e-3, + "scale": 1.0, + "damping": 1.0, + "use_grad": False, + "constraint_lr": 0.0, + } + } + + +def _conv_weight(c_out=4, c_in=2, kh=3, kw=3, seed=42): + """Torch conv kernel: OIHW (C_out, C_in, kH, kW).""" + rng = np.random.default_rng(seed) + return torch.tensor(rng.standard_normal((c_out, c_in, kh, kw)).astype(np.float32)) + + +# --- Metrics --- + + +def test_unstructured_sparsity_value(): + """weight=[0,0,0,2], target=0.5 -> L0=0.75, factor=-0.3125, metric=-0.15625.""" + metric = UnstructuredSparsityMetric(l0_mode="coarse", scale_mode="mean", epsilon=1e-3, target_sparsity=0.5) + assert abs(float(metric(torch.tensor([0.0, 0.0, 0.0, 2.0]))) - (-0.15625)) < 1e-5 + + +@pytest.mark.parametrize( + "target_resource,fill,expected", + [("DSP", 0.0, 1.0), ("DSP", 1.0, 0.0), ("BRAM", 0.0, 1.0), ("BRAM", 1.0, 0.0)], +) +def test_fpga_extreme_weights(target_resource, fill, expected): + metric = FPGAAwareSparsityMetric(rf=2, precision=16, bram_width=36, target_resource=target_resource, epsilon=1e-3) + assert abs(float(metric(torch.full((4, 16), fill))) - expected) < 1e-5 + + +def test_fpga_dsp_l2_grouping_math(): + metric = FPGAAwareSparsityMetric(rf=2, target_resource="DSP", epsilon=1e-3) + assert abs(float(metric(torch.tensor([[1.0, 1.0, 0.0, 0.0]]))) - 0.5) < 1e-5 + + +@pytest.mark.parametrize("precision,bram_width,expected_c", [(16, 36, 4), (18, 36, 2)]) +def test_fpga_calculate_c(precision, bram_width, expected_c): + assert FPGAAwareSparsityMetric(precision=precision, bram_width=bram_width, target_resource="DSP").c == expected_c + + +def test_fpga_handles_1d_weight(): + metric = FPGAAwareSparsityMetric(rf=2, target_resource="DSP", epsilon=1e-3) + assert abs(float(metric(torch.ones(8)))) < 1e-5 + + +def test_fpga_invalid_target_resource_raises_on_call(): + with pytest.raises(ValueError): + FPGAAwareSparsityMetric(target_resource="LUT")(torch.ones(4, 8)) + + +def test_fpga_bram_packing_impossible_raises_on_call(): + with pytest.raises(ValueError): + FPGAAwareSparsityMetric(precision=128, bram_width=16, target_resource="BRAM")(torch.ones(4, 16)) + + +def test_paca_returns_zero_for_non_4d(): + assert abs(float(PACAPatternMetric()(torch.rand(8, 4)))) < 1e-6 + + +def test_paca_lazy_caching(): + metric = PACAPatternMetric(num_patterns_to_keep=4, beta=0.75) + w = _conv_weight() + assert metric.dominant_patterns is None + _ = metric(w) + first = metric.dominant_patterns + assert first is not None + _ = metric(w) + assert metric.dominant_patterns is first + + +def test_paca_projection_mask_shape_and_binary(): + metric = PACAPatternMetric(num_patterns_to_keep=4, beta=0.75) + w = _conv_weight() + _ = metric(w) + mask = metric.get_projection_mask(w) + assert tuple(mask.shape) == tuple(w.shape) + assert set(np.unique(mask.detach().cpu().numpy()).tolist()).issubset({0, 1}) + + +@pytest.mark.parametrize("distance_metric", ["hamming", "valued_hamming", "cosine"]) +def test_paca_all_distance_metrics(distance_metric): + metric = PACAPatternMetric(num_patterns_to_keep=4, beta=0.75, distance_metric=distance_metric) + assert math.isfinite(float(metric(_conv_weight()))) + + +def test_paca_projection_mask_identity_when_no_patterns(): + metric = PACAPatternMetric(num_patterns_to_keep=4, beta=0.75) + w = _conv_weight() + np.testing.assert_array_equal( + metric.get_projection_mask(w).detach().cpu().numpy(), np.ones(tuple(w.shape), dtype=np.float32) + ) + + +def test_paca_invalid_distance_metric_raises_on_call(): + with pytest.raises(ValueError): + PACAPatternMetric(distance_metric="manhattan")(_conv_weight()) + + +# --- Pattern helpers --- + + +def test_patterns_kernels_and_patterns(): + weight = torch.tensor([[[[1.0, 0.0], [0.0, 1.0]]]]) # (1,1,2,2) OIHW + kernels, pats, _ = patterns.kernels_and_patterns(weight, src="OIHW", epsilon=0.5) + assert tuple(kernels.shape) == (1, 4) + np.testing.assert_array_equal(pats.detach().cpu().numpy(), np.array([[1, 0, 0, 1]], dtype=np.uint8)) + + +def test_patterns_select_dominant_fixed_size_with_validity(): + pats = torch.tensor(np.array([[1, 1, 0, 0]] * 6 + [[0, 0, 1, 1]] * 2, dtype=np.uint8)) + dom, valid = patterns.select_dominant_patterns(pats, num_patterns_to_keep=4, beta=0.99) + assert tuple(dom.shape) == (4, 4) + valid_np = valid.detach().cpu().numpy() + assert valid_np[:2].all() and not valid_np[2:].any() + + +# --- MDMM phase cycle --- + + +def test_mdmm_get_layer_sparsity(): + weight = torch.tensor([[1e-4, 1e-3, 1e-2, 0.1], [-1e-4, -1e-3, -1e-2, -0.1]]) + mdmm = MDMM(base_config(), "linear") + mdmm.build(weight.shape) + assert abs(float(mdmm.get_layer_sparsity(weight)) - 0.5) < 1e-5 + + +def test_mdmm_fpga_phase_cycle(): + cfg = base_config() + cfg["pruning_parameters"].update( + {"metric_type": "FPGAAwareSparsity", "precision": 16, "target_resource": "DSP", "rf": 2} + ) + weight = _conv_weight() + mdmm = MDMM(cfg, "conv") + mdmm.build(weight.shape) + assert torch.equal(mdmm(weight), weight) # pretraining: unchanged + assert abs(float(mdmm.calculate_additional_loss())) < 1e-8 + mdmm.post_pre_train_function() + _ = mdmm(weight) # active + mdmm.pre_finetune_function() + _ = mdmm(weight) # finetuning + assert abs(float(mdmm.calculate_additional_loss())) < 1e-8 + + +def test_mdmm_paca_forces_equality(): + cfg = base_config() + cfg["pruning_parameters"].update( + {"metric_type": "PACAPatternSparsity", "constraint_type": "GreaterThanOrEqual", "target_value": 99.0, + "num_patterns_to_keep": 4, "beta": 0.85, "distance_metric": "cosine"} + ) + mdmm = MDMM(cfg, "conv") + mdmm.build(_conv_weight().shape) + assert isinstance(mdmm.constraint_layer, EqualityConstraint) + assert mdmm.constraint_layer.target_value == 0.0 + + +def test_mdmm_paca_phase_cycle_projection_mask(): + cfg = base_config() + cfg["pruning_parameters"].update( + {"metric_type": "PACAPatternSparsity", "num_patterns_to_keep": 4, "beta": 0.85, "distance_metric": "cosine"} + ) + weight = _conv_weight(c_out=4, c_in=4) + mdmm = MDMM(cfg, "conv") + mdmm.build(weight.shape) + assert torch.equal(mdmm(weight), weight) # pretraining + assert mdmm.constraint_layer.metric_fn.dominant_patterns is not None + mdmm.post_pre_train_function() + _ = mdmm(weight) # active + mdmm.pre_finetune_function() + out = mdmm(weight) # finetuning -> weight * binary mask + out_np, weight_np = out.detach().cpu().numpy(), weight.detach().cpu().numpy() + assert np.logical_or(np.isclose(out_np, 0.0), np.isclose(out_np, weight_np)).all()