apple / apple/coreai-optimization

[Bug]: KMeansPalettizer leaks checkpoint files and leaves fake palettization disabled on calibration failure; fails to resolve root module sensitivities

Open
#94 0 comments 0 reactions 2 assignees Claimed by @crowbat View on GitHub
Dominant language
Python
Stars
122
Forks
36
Avg merge
20h 50m
Merged PRs (30d)
17

Description

### Describe the Bug

Two related issues in `KMeansPalettizer` cause resource leaks, state corruption, and silent sensitivity configuration failures during calibration:

1. **Resource Leak & State Corruption on Calibration Abort (`src/coreai_opt/palettization/kmeans/palettizer.py:285-340`)**:
When entering `KMeansPalettizer.calibration_mode()`, the palettizer saves a temporary model checkpoint to disk via `tempfile.NamedTemporaryFile(delete=False, suffix=".pt", prefix="palettizer_calibration_")`.
The deletion `os.unlink(checkpoint_path)` only takes place inside `_load_model_checkpoint()`.
However, the post-calibration logic was located inside the generator's `finally:` block. Consequently:
- If an exception is raised inside the user's `with palettizer.calibration_mode(...)` block prior to calling `skm.step()` (e.g. data loader error, OOM during forward pass, or user assertion), the `finally:` block executes `if not calibration_helper.step_called: raise RuntimeError(...)`. In Python, raising in a `finally:` block **completely masks and suppresses the caller's original exception**.
- Because `_load_model_checkpoint()` is never reached, the multi-gigabyte `.pt` model checkpoint file is **permanently leaked in `/tmp`**. On distributed cluster nodes or shared servers, repeated failed calibration attempts can quickly cause disk exhaustion.
- `self._model.apply(_disable_fake_palett)` was called at context entry, but `self._model.apply(_enable_fake_palett)` is never reached on abort. The model is left permanently mutated with fake palettization silently disabled.
- Model weights modified during calibration are left un-restored.

2. **Root Module Sensitivity Parameter Name Mismatch (`src/coreai_opt/palettization/kmeans/palettizer.py:685, 767`)**:
`_set_sensitivities_in_fake_palettize_modules` and `_get_sensitivities_from_fake_palettize_modules` build parameter names using:
```python
param_name = ".".join([module_name, "parametrizations", attr_name, "original"])
```
When palettizing a standalone root module (e.g., `model = nn.Linear(...)`), `module_name == ""`. This produces `".parametrizations.weight.original"` with an invalid leading dot.
Because PyTorch's `self._model.named_parameters()` produces `"parametrizations.weight.original"` (without a leading dot), `param_name not in sensitivity_dict` triggers every time, logging:
```text
ERROR: No sensitivity value found for .parametrizations.weight.original
```
Sensitivities are therefore never assigned to root-level modules.

---

### Steps to Reproduce

#### Case A: Resource leak & exception masking on calibration failure
```python
import glob, os, tempfile
import torch, torch.nn as nn
from coreai_opt.palettization import KMeansPalettizer, KMeansPalettizerConfig, ModuleKMeansPalettizerConfig
from coreai_opt.palettization.spec import default_weight_palettization_spec

model = nn.Sequential(nn.Linear(4, 4))
config = KMeansPalettizerConfig(
global_config=ModuleKMeansPalettizerConfig(
op_state_spec={"weight": default_weight_palettization_spec()}
)
)
palettizer = KMeansPalettizer(model, config)
prepared = palettizer.prepare((torch.randn(2, 4),))

tmp_dir = tempfile.gettempdir()
before = set(glob.glob(os.path.join(tmp_dir, "palettizer_calibration_*.pt")))

class DataLoadingError(Exception): pass

try:
with palettizer.calibration_mode(loss_fn=lambda out, tgt: out.sum()) as skm:
raise DataLoadingError("DataLoader failed before step")
except BaseException as e:
print("Caught:", type(e).__name__, str(e))

after = set(glob.glob(os.path.join(tmp_dir, "palettizer_calibration_*.pt")))
print("Leaked checkpoints count:", len(after - before))
```
**Output on `main`**:
```text
Caught: RuntimeError calibration_mode requires at least one call to step(). No calibration data was processed.
Leaked checkpoints count: 1
```
*(Notice: `DataLoadingError` was masked, and a `.pt` file was leaked).*

#### Case B: Root module sensitivity failure
```python
model = nn.Linear(4, 4)
palettizer = KMeansPalettizer(model, config)
prepared = palettizer.prepare((torch.randn(2, 4),))

with palettizer.calibration_mode(loss_fn=lambda out, tgt: out.sum()) as skm:
out = prepared(torch.randn(2, 4))
skm.step(out, torch.zeros(2, 4))
```
**Output on `main`**:
```text
ERROR coreai_opt.palettization.kmeans.palettizer: No sensitivity value found for .parametrizations.weight.original
```

---

### Expected Behavior
1. If an error occurs inside `calibration_mode`:
- The caller's original exception should propagate without being masked.
- Temporary `.pt` checkpoint files in `/tmp` must be deterministically unlinked.
- Model weights must be rolled back to their pre-calibration state.
- Fake palettization must be re-enabled.
- Palettizer lifecycle must reset to `IDLE`.
2. Root modules must resolve parameter names canonically (`"parametrizations.weight.original"`) without leading dots, correctly attaching sensitivity tensors.

---

### Environment
- OS: macOS (Darwin)
- Python: 3.11
- PyTorch: 2.13.0
- Commit: `main` (73fd849)

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.