Binary optimizers: inference, fit suite, and training infrastructure#1
Binary optimizers: inference, fit suite, and training infrastructure#1leo-gan wants to merge 9 commits into
Conversation
…ipeline. Introduce XNOR+popcount binary linear layers with weight packing, training/inference memory profiling, optimizer×model training sweeps, Pareto analysis, experiments, tests, and scaffold results docs.
Scaffold empty cosine_voting, ema_flip, hybrid_accumulator, and sparse_sign optimizers; include test_debug.py for environment checks.
…timizers. Fill previously empty modules with binary-weight optimizers (cosine-annealed voting, EMA-gated flips, adaptive integrate-and-fire hybrid, sparse sign updates) and export them from the optimizers package.
Integrate EMAFlip, CosineVoting, SparseSign, and HybridAccumulator into the training sweep, fix CosineVoting state handling, add improvement-report generation and tests, and refresh scaffold Pareto/summary docs.
Correct Hybrid fire to set weights to sign(acc), restore EMAFlip adaptive continuous updates, raise CosineVoting LR defaults, update sweep hyperparams, and add regression tests.
…izers. Fingerprint-based save/load under checkpoints/ (gitignored); EMAFlip cosine anneal and Hybrid soft fire; SparseSign density annealing; CosineVoting restarts; fit and benchmark CLI entrypoints.
Add HybridV2 (momentum + cosine LR + sparse + soft IF), DualOptimizer for BN/bias Adam, schema fit_v3 fingerprints, and refreshed fit/benchmark reports. Checkpoints remain local-only (gitignored).
Merge optimizers.md with NEW_OPTIMIZERS into one design/trade-off guide, document fit_v3 rationale for the default, add improvement proposals, and update fit CLI/reports to drop old-vs-new framing.
… suite. Wire swarm_mlp with swarm/swarm_log/adam into fit comparisons, improve SparseSign density hold and long-run docs, re-enable research optimizers with potential while pausing hybrid_v2, and refresh fit reports.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a comprehensive benchmarking, profiling, and inference infrastructure for binary neural networks, including a bitpacked inference engine, memory profilers, and several new binary optimizers (such as ema_flip and cosine_voting). The reviewer's feedback highlights key improvement opportunities: dynamically merging optimizer states in DualOptimizer to ensure accurate memory profiling of continuous parameters, wrapping warm-up steps in a try-except block to prevent crashes on empty data loaders, and avoiding direct .data modifications in CosineVotingOptimizer by operating directly on tensors under @torch.no_grad().
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| # Expose state for memory profilers that inspect .state | ||
| self.state = binary_opt.state | ||
| self.param_groups = list(binary_opt.param_groups) | ||
| if bn_opt is not None: | ||
| self.param_groups = list(binary_opt.param_groups) + list(bn_opt.param_groups) |
There was a problem hiding this comment.
Exposing self.state = binary_opt.state in __init__ only captures the state of the binary optimizer. Since DualOptimizer also manages bn_opt (which is typically a memory-heavy optimizer like Adam), the memory profiler will completely miss the continuous parameters' optimizer state. Implementing state as a dynamic property that merges both states resolves this profiling discrepancy.
| # Expose state for memory profilers that inspect .state | |
| self.state = binary_opt.state | |
| self.param_groups = list(binary_opt.param_groups) | |
| if bn_opt is not None: | |
| self.param_groups = list(binary_opt.param_groups) + list(bn_opt.param_groups) | |
| self.param_groups = list(binary_opt.param_groups) | |
| if bn_opt is not None: | |
| self.param_groups = list(binary_opt.param_groups) + list(bn_opt.param_groups) | |
| @property | |
| def state(self): | |
| combined = dict(self.binary_opt.state) | |
| if self.bn_opt is not None: | |
| combined.update(self.bn_opt.state) | |
| return combined |
| # Warm-up for memory profile | ||
| model.train() | ||
| xb, yb = next(iter(train_loader)) | ||
| xb, yb = xb.to(device), yb.to(device) | ||
| optimizer.zero_grad() | ||
| model(xb).sum().backward() | ||
| optimizer.step() | ||
| optimizer.zero_grad() |
There was a problem hiding this comment.
In train_or_load, calling next(iter(train_loader)) directly can raise a StopIteration exception if the training loader is empty or has no batches. To prevent training from crashing, wrap the warm-up step in a try-except StopIteration block, similar to the implementation in training_sweep.py.
| # Warm-up for memory profile | |
| model.train() | |
| xb, yb = next(iter(train_loader)) | |
| xb, yb = xb.to(device), yb.to(device) | |
| optimizer.zero_grad() | |
| model(xb).sum().backward() | |
| optimizer.step() | |
| optimizer.zero_grad() | |
| # Warm-up for memory profile | |
| model.train() | |
| try: | |
| xb, yb = next(iter(train_loader)) | |
| xb, yb = xb.to(device), yb.to(device) | |
| optimizer.zero_grad() | |
| model(xb).sum().backward() | |
| optimizer.step() | |
| optimizer.zero_grad() | |
| except StopIteration: | |
| pass |
| if p.data.dim() < 2: | ||
| p.data.add_(p.grad.data, alpha=-bn_lr) | ||
| continue | ||
|
|
||
| state = self.state[p] | ||
| if "momentum_buffer" not in state: | ||
| state["momentum_buffer"] = torch.zeros_like(p.data) | ||
|
|
||
| buf = state["momentum_buffer"] | ||
| buf.mul_(mom).add_(p.grad.data, alpha=1.0 - mom) | ||
|
|
||
| if conf_thresh > 0: | ||
| vote = torch.zeros_like(buf) | ||
| confident = buf.abs() > conf_thresh | ||
| vote[confident] = torch.sign(buf[confident]) | ||
| else: | ||
| vote = torch.sign(buf) | ||
| # sign(0) → 0; treat zeros as no-op already | ||
|
|
||
| p.data.add_(vote, alpha=-lr) | ||
| p.data.clamp_(-clip, clip) |
There was a problem hiding this comment.
Directly accessing and modifying .data (e.g., p.data.add_, p.grad.data) is discouraged in modern PyTorch as it can bypass autograd tracking, version counters, and cause silent bugs. Since the step method is already decorated with @torch.no_grad(), you can perform these operations directly on the parameter tensors.
if p.dim() < 2:
p.add_(p.grad, alpha=-bn_lr)
continue
state = self.state[p]
if "momentum_buffer" not in state:
state["momentum_buffer"] = torch.zeros_like(p)
buf = state["momentum_buffer"]
buf.mul_(mom).add_(p.grad, alpha=1.0 - mom)
if conf_thresh > 0:
vote = torch.zeros_like(buf)
confident = buf.abs() > conf_thresh
vote[confident] = torch.sign(buf[confident])
else:
vote = torch.sign(buf)
# sign(0) → 0; treat zeros as no-op already
p.add_(vote, alpha=-lr)
p.clamp_(-clip, clip)1 participant
Summary
checkpoints/(local only, gitignored), fit + re-benchmark CLIs, Bit-MLP suite plus swarm_mlp × (swarm, swarm_log, swarm_log_dynamic, adam) comparisons.docs/optimizers.md(reasons/trade-offs, active vs paused), fit analysis, improvement proposals, SparseSign max-fit notes.Key results (fit protocol, MNIST)
ema_flipbest among STE FC trainers (~96.3% best test @ 15 ep); SparseSign reaches ~97.1% with 40-ep dense-hold recipe.swarm_log_dynamicand Adam-on-swarm competitive; classic recruit-rateswarmlower (~90%).Test plan
uv run pytest tests/ -q(packing, binary linear, memory, optimizers, checkpoints, Pareto)uv run python experiments/run_fit_training.py --default-only(cache hit or short train)uv run python experiments/run_fit_training.py --swarm-onlyif swarm comparison neededuv run python experiments/run_benchmark_checkpoints.py(eval only, no train)uv run python experiments/run_full_pipeline.py(scaffold, not fit-scale)