torch-harness 0.10.1__tar.gz

torch_harness-0.10.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Vadym Stupakov
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

torch_harness-0.10.1/PKG-INFO

@@ -0,0 +1,236 @@
+Metadata-Version: 2.4
+Name: torch-harness
+Version: 0.10.1
+Summary: Reliable utilities for PyTorch and Lightning
+Keywords: lightning,machine-learning,pytorch,torch
+Author: Vadym Stupakov
+Author-email: Vadym Stupakov <vadim.stupakov@gmail.com>
+License-Expression: MIT
+License-File: LICENSE
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Typing :: Typed
+Requires-Dist: jaxtyping>=0.3.11
+Requires-Dist: loguru>=0.7
+Requires-Dist: torch>=2.0
+Requires-Dist: typing-extensions>=4.4
+Requires-Dist: lightning>=2.6,<3 ; extra == 'lightning'
+Requires-Dist: onnx>=1.18 ; extra == 'onnx'
+Requires-Python: >=3.11
+Project-URL: Repository, https://github.com/Red-Eyed/torch-harness
+Project-URL: Issues, https://github.com/Red-Eyed/torch-harness/issues
+Provides-Extra: lightning
+Provides-Extra: onnx
+Description-Content-Type: text/markdown
+
+# torch-harness
+
+An intentionally uncategorized collection of utilities for PyTorch modeling
+experiments. It is a toolbox rather than a single framework: utilities are added
+as experimental needs arise, without forcing them into an artificial hierarchy.
+
+Each utility is modular and self-contained. It owns its implementation and tests,
+avoids assumptions about the surrounding project, and can be used independently
+of the other utilities. Optional framework integrations remain isolated so they
+do not add dependencies to the core package.
+
+## Adaptive loss balancing
+
+`LossBalancer` adapts named auxiliary weights from shared gradient norms while
+keeping the primary loss weight exactly one:
+
+```python
+from torch_harness.losses import LossBalancer, LossTarget, weighted_loss
+
+balancer = LossBalancer(
+    reference="primary",
+    adaptive={"auxiliary": LossTarget(ratio=0.3, initial_weight=0.1)},
+)
+weights = balancer.update(losses, shared=shared_features, step=global_step)
+total_loss = weighted_loss(losses, weights)
+```
+
+It supports independent targets, fixed weights, optional losses, periodic EMA
+updates, checkpoint continuation, and synchronized DDP activation measurements.
+See the [component documentation](src/torch_harness/losses/adaptive_balancing/README.md)
+for stability controls, diagnostics, and training integration. Run
+`uv run examples/adaptive_loss_balancing.py` for a synthetic optimization example.
+
+## Uncertainty weighting
+
+`UncertaintyWeighting` learns positive weights for a vector of peer task losses:
+
+```python
+import torch
+
+from torch_harness.losses import UncertaintyWeighting
+
+weighting = UncertaintyWeighting(num_tasks=3)
+task_losses = torch.stack((classification_loss, depth_loss, normal_loss))
+loss = weighting(task_losses)
+weights_for_logging = weighting.effective_weights
+```
+
+Each effective weight starts near one and adapts jointly with the model. The
+module preserves gradients to the input task losses and its uncertainty
+parameters, while the weights exposed for logging are detached. See the
+[component documentation](src/torch_harness/losses/uncertainty_weighting/README.md)
+for the exact objective, optimizer wiring, numerical behavior, and limitations.
+
+## SuperLoss
+
+`SuperLoss` applies robust curriculum weighting to any unreduced task loss:
+
+```python
+from math import log
+
+import torch.nn.functional as F
+
+from torch_harness.losses import SuperLoss
+
+criterion = SuperLoss(threshold=log(10), regularization=1.0)
+task_loss = F.cross_entropy(logits, targets, reduction="none")
+loss = criterion(task_loss)
+```
+
+It computes the closed-form optimal sample confidence from the SuperLoss paper,
+upweighting easy samples and downweighting hard samples. The implementation is
+task-agnostic, runs entirely in PyTorch, and has no SciPy or per-sample-state
+dependency. See the [component documentation](src/torch_harness/losses/super_loss/README.md)
+for parameter guidance and numerical details.
+
+## SafeBatchNorm
+
+`SafeBatchNorm` rejects a non-finite activation before it can corrupt BatchNorm
+running statistics.
+
+```python
+from torch_harness.layers import SafeBatchNorm2d
+
+normalization = SafeBatchNorm2d(64)
+```
+
+`SafeBatchNorm1d`, `SafeBatchNorm2d`, and `SafeBatchNorm3d` directly inherit
+their matching PyTorch classes. They retain the native constructor, state-dict
+layout, and type identity. A non-finite input raises `FloatingPointError`, which
+a fault-tolerant training loop can catch to skip the step.
+
+`nn.SyncBatchNorm` is intentionally unsupported because a rank-local failure
+before its collective could deadlock the other ranks.
+
+## Runtime layer replacement
+
+`replace_layers` recursively transforms existing models through an explicit
+replacement factory:
+
+```python
+from torch import nn
+
+from torch_harness.runtime import replace_layers
+
+model = replace_layers(
+    model,
+    old_layer_cls=nn.SiLU,
+    replacement_factory=lambda silu: nn.ReLU(inplace=silu.inplace),
+)
+```
+
+The traversal includes nested containers such as `Sequential`. Run replacement
+before constructing the optimizer, distributed wrappers, or a compiled model.
+
+## Fault-tolerant training steps
+
+`FaultTolerantTrainingStep` skips occasional exceptions and NaN or infinite losses
+before Lightning automatic optimization runs backward or updates the optimizer:
+
+```python
+from datetime import timedelta
+
+from lightning.pytorch import Trainer
+from torch_harness.lightning import FaultTolerantTrainingStep
+
+fault_tolerance = FaultTolerantTrainingStep(
+    max_faults=3,
+    fault_window=timedelta(hours=1),
+)
+trainer = Trainer(callbacks=[fault_tolerance])
+```
+
+This example skips the first three faults in any rolling one-hour window. A fourth
+fault within that window is logged through Loguru and propagated. Exceptions retain
+their original type and traceback; repeated non-finite losses raise
+`NonFiniteLossError`.
+
+Install the Lightning integration with `torch-harness[lightning]`.
+
+The callback intentionally supports only single-process automatic optimization.
+Errors raised during backward or `optimizer.step` happen after the recoverable
+boundary and are propagated because an optimizer update may already be partial.
+
+## Model structures in files
+
+`FileModelStructure` saves the recursive structure produced by `print(model)` as
+readable UTF-8 text when fitting starts:
+
+```python
+from pathlib import Path
+
+from lightning.pytorch import Trainer
+from torch_harness.lightning import FileModelStructure
+
+structure = FileModelStructure(
+    output_path=Path("artifacts/model-structure.txt"),
+)
+trainer = Trainer(callbacks=[structure])
+```
+
+The callback includes the complete registered module hierarchy without requiring
+example inputs or running a forward pass. It creates parent directories, replaces
+an existing structure file, and writes only from the global-zero process. Runtime
+tensor operations that are not registered modules do not appear.
+
+Run `uv run examples/file_model_structure.py` to generate an inspectable example
+at `examples/model-structure.txt`.
+
+## Human-readable ONNX graphs
+
+Install `torch-harness[onnx]` to save module summaries suitable for Git diffs
+and Meld:
+
+```python
+from pathlib import Path
+
+import onnx
+
+from torch_harness.onnx import write_text_graph
+
+model = onnx.load("model.onnx")
+write_text_graph(model, Path("artifacts/model.onnx.txt"))
+```
+
+Summaries show operation signatures and counts, attributes, small constants,
+parameter shapes, and input/output dimensions. They omit temporary node indexes
+and weight payloads, making RMSNorm decomposition versus `RMSNormalization`
+visible without renumbering unrelated operations. Internal wiring is omitted.
+
+See the [ONNX documentation](src/torch_harness/onnx/README.md) for source module
+capture, shape inference, and format limitations. Run
+`uv run --extra onnx examples/onnx_text_graph.py` for a small example.
+
+## Planned utilities
+
+- A Lightning `PreciseBN` callback for recomputing BatchNorm running statistics.
+- A Lightning mixin that manages schedule-free AdamW train/evaluation state.
+
+Each integration will be isolated so users only install the frameworks they need.
+
+## License
+
+`torch-harness` is released under the MIT License.

torch_harness-0.10.1/README.md

@@ -0,0 +1,204 @@
+# torch-harness
+
+An intentionally uncategorized collection of utilities for PyTorch modeling
+experiments. It is a toolbox rather than a single framework: utilities are added
+as experimental needs arise, without forcing them into an artificial hierarchy.
+
+Each utility is modular and self-contained. It owns its implementation and tests,
+avoids assumptions about the surrounding project, and can be used independently
+of the other utilities. Optional framework integrations remain isolated so they
+do not add dependencies to the core package.
+
+## Adaptive loss balancing
+
+`LossBalancer` adapts named auxiliary weights from shared gradient norms while
+keeping the primary loss weight exactly one:
+
+```python
+from torch_harness.losses import LossBalancer, LossTarget, weighted_loss
+
+balancer = LossBalancer(
+    reference="primary",
+    adaptive={"auxiliary": LossTarget(ratio=0.3, initial_weight=0.1)},
+)
+weights = balancer.update(losses, shared=shared_features, step=global_step)
+total_loss = weighted_loss(losses, weights)
+```
+
+It supports independent targets, fixed weights, optional losses, periodic EMA
+updates, checkpoint continuation, and synchronized DDP activation measurements.
+See the [component documentation](src/torch_harness/losses/adaptive_balancing/README.md)
+for stability controls, diagnostics, and training integration. Run
+`uv run examples/adaptive_loss_balancing.py` for a synthetic optimization example.
+
+## Uncertainty weighting
+
+`UncertaintyWeighting` learns positive weights for a vector of peer task losses:
+
+```python
+import torch
+
+from torch_harness.losses import UncertaintyWeighting
+
+weighting = UncertaintyWeighting(num_tasks=3)
+task_losses = torch.stack((classification_loss, depth_loss, normal_loss))
+loss = weighting(task_losses)
+weights_for_logging = weighting.effective_weights
+```
+
+Each effective weight starts near one and adapts jointly with the model. The
+module preserves gradients to the input task losses and its uncertainty
+parameters, while the weights exposed for logging are detached. See the
+[component documentation](src/torch_harness/losses/uncertainty_weighting/README.md)
+for the exact objective, optimizer wiring, numerical behavior, and limitations.
+
+## SuperLoss
+
+`SuperLoss` applies robust curriculum weighting to any unreduced task loss:
+
+```python
+from math import log
+
+import torch.nn.functional as F
+
+from torch_harness.losses import SuperLoss
+
+criterion = SuperLoss(threshold=log(10), regularization=1.0)
+task_loss = F.cross_entropy(logits, targets, reduction="none")
+loss = criterion(task_loss)
+```
+
+It computes the closed-form optimal sample confidence from the SuperLoss paper,
+upweighting easy samples and downweighting hard samples. The implementation is
+task-agnostic, runs entirely in PyTorch, and has no SciPy or per-sample-state
+dependency. See the [component documentation](src/torch_harness/losses/super_loss/README.md)
+for parameter guidance and numerical details.
+
+## SafeBatchNorm
+
+`SafeBatchNorm` rejects a non-finite activation before it can corrupt BatchNorm
+running statistics.
+
+```python
+from torch_harness.layers import SafeBatchNorm2d
+
+normalization = SafeBatchNorm2d(64)
+```
+
+`SafeBatchNorm1d`, `SafeBatchNorm2d`, and `SafeBatchNorm3d` directly inherit
+their matching PyTorch classes. They retain the native constructor, state-dict
+layout, and type identity. A non-finite input raises `FloatingPointError`, which
+a fault-tolerant training loop can catch to skip the step.
+
+`nn.SyncBatchNorm` is intentionally unsupported because a rank-local failure
+before its collective could deadlock the other ranks.
+
+## Runtime layer replacement
+
+`replace_layers` recursively transforms existing models through an explicit
+replacement factory:
+
+```python
+from torch import nn
+
+from torch_harness.runtime import replace_layers
+
+model = replace_layers(
+    model,
+    old_layer_cls=nn.SiLU,
+    replacement_factory=lambda silu: nn.ReLU(inplace=silu.inplace),
+)
+```
+
+The traversal includes nested containers such as `Sequential`. Run replacement
+before constructing the optimizer, distributed wrappers, or a compiled model.
+
+## Fault-tolerant training steps
+
+`FaultTolerantTrainingStep` skips occasional exceptions and NaN or infinite losses
+before Lightning automatic optimization runs backward or updates the optimizer:
+
+```python
+from datetime import timedelta
+
+from lightning.pytorch import Trainer
+from torch_harness.lightning import FaultTolerantTrainingStep
+
+fault_tolerance = FaultTolerantTrainingStep(
+    max_faults=3,
+    fault_window=timedelta(hours=1),
+)
+trainer = Trainer(callbacks=[fault_tolerance])
+```
+
+This example skips the first three faults in any rolling one-hour window. A fourth
+fault within that window is logged through Loguru and propagated. Exceptions retain
+their original type and traceback; repeated non-finite losses raise
+`NonFiniteLossError`.
+
+Install the Lightning integration with `torch-harness[lightning]`.
+
+The callback intentionally supports only single-process automatic optimization.
+Errors raised during backward or `optimizer.step` happen after the recoverable
+boundary and are propagated because an optimizer update may already be partial.
+
+## Model structures in files
+
+`FileModelStructure` saves the recursive structure produced by `print(model)` as
+readable UTF-8 text when fitting starts:
+
+```python
+from pathlib import Path
+
+from lightning.pytorch import Trainer
+from torch_harness.lightning import FileModelStructure
+
+structure = FileModelStructure(
+    output_path=Path("artifacts/model-structure.txt"),
+)
+trainer = Trainer(callbacks=[structure])
+```
+
+The callback includes the complete registered module hierarchy without requiring
+example inputs or running a forward pass. It creates parent directories, replaces
+an existing structure file, and writes only from the global-zero process. Runtime
+tensor operations that are not registered modules do not appear.
+
+Run `uv run examples/file_model_structure.py` to generate an inspectable example
+at `examples/model-structure.txt`.
+
+## Human-readable ONNX graphs
+
+Install `torch-harness[onnx]` to save module summaries suitable for Git diffs
+and Meld:
+
+```python
+from pathlib import Path
+
+import onnx
+
+from torch_harness.onnx import write_text_graph
+
+model = onnx.load("model.onnx")
+write_text_graph(model, Path("artifacts/model.onnx.txt"))
+```
+
+Summaries show operation signatures and counts, attributes, small constants,
+parameter shapes, and input/output dimensions. They omit temporary node indexes
+and weight payloads, making RMSNorm decomposition versus `RMSNormalization`
+visible without renumbering unrelated operations. Internal wiring is omitted.
+
+See the [ONNX documentation](src/torch_harness/onnx/README.md) for source module
+capture, shape inference, and format limitations. Run
+`uv run --extra onnx examples/onnx_text_graph.py` for a small example.
+
+## Planned utilities
+
+- A Lightning `PreciseBN` callback for recomputing BatchNorm running statistics.
+- A Lightning mixin that manages schedule-free AdamW train/evaluation state.
+
+Each integration will be isolated so users only install the frameworks they need.
+
+## License
+
+`torch-harness` is released under the MIT License.

torch_harness-0.10.1/pyproject.toml

@@ -0,0 +1,58 @@
+[project]
+name = "torch-harness"
+version = "0.10.1"
+description = "Reliable utilities for PyTorch and Lightning"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "jaxtyping>=0.3.11",
+    "loguru>=0.7",
+    "torch>=2.0",
+    "typing-extensions>=4.4",
+]
+license = "MIT"
+license-files = ["LICENSE"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
+    "Typing :: Typed",
+]
+keywords = [
+    "lightning",
+    "machine-learning",
+    "pytorch",
+    "torch",
+]
+
+[[project.authors]]
+name = "Vadym Stupakov"
+email = "vadim.stupakov@gmail.com"
+
+[project.optional-dependencies]
+lightning = ["lightning>=2.6,<3"]
+onnx = ["onnx>=1.18"]
+
+[project.urls]
+Repository = "https://github.com/Red-Eyed/torch-harness"
+Issues = "https://github.com/Red-Eyed/torch-harness/issues"
+
+[build-system]
+requires = ["uv_build>=0.12.1,<0.13.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "lightning>=2.6,<3",
+    "onnx>=1.18",
+    "pyrefly>=1.2.0",
+    "pytest>=9.1.1",
+    "ruff>=0.16.4",
+    "scipy>=1.17.1",
+]

torch_harness-0.10.1/pyproject.toml.orig

@@ -0,0 +1,52 @@
+[project]
+name = "torch-harness"
+version = "0.10.1"
+description = "Reliable utilities for PyTorch and Lightning"
+readme = "README.md"
+authors = [
+    { name = "Vadym Stupakov", email = "vadim.stupakov@gmail.com" },
+]
+requires-python = ">=3.11"
+dependencies = [
+    "jaxtyping>=0.3.11",
+    "loguru>=0.7",
+    "torch>=2.0",
+    "typing-extensions>=4.4",
+]
+license = "MIT"
+license-files = ["LICENSE"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Programming Language :: Python :: 3.14",
+    "Typing :: Typed",
+]
+keywords = ["lightning", "machine-learning", "pytorch", "torch"]
+
+[project.optional-dependencies]
+lightning = ["lightning>=2.6,<3"]
+onnx = ["onnx>=1.18"]
+
+[project.urls]
+Repository = "https://github.com/Red-Eyed/torch-harness"
+Issues = "https://github.com/Red-Eyed/torch-harness/issues"
+
+[build-system]
+requires = ["uv_build>=0.12.1,<0.13.0"]
+build-backend = "uv_build"
+
+[dependency-groups]
+dev = [
+    "lightning>=2.6,<3",
+    "onnx>=1.18",
+    "pyrefly>=1.2.0",
+    "pytest>=9.1.1",
+    "ruff>=0.16.4",
+    "scipy>=1.17.1",
+]

torch_harness-0.10.1/src/torch_harness/__init__.py

@@ -0,0 +1,27 @@
+"""Reliable utilities for PyTorch and Lightning."""
+
+from torch_harness.layers import (
+    SafeBatchNorm1d,
+    SafeBatchNorm2d,
+    SafeBatchNorm3d,
+)
+from torch_harness.losses import (
+    BalancerConfig,
+    LossBalancer,
+    LossTarget,
+    SuperLoss,
+    UncertaintyWeighting,
+    weighted_loss,
+)
+
+__all__ = [
+    "BalancerConfig",
+    "LossBalancer",
+    "LossTarget",
+    "SafeBatchNorm1d",
+    "SafeBatchNorm2d",
+    "SafeBatchNorm3d",
+    "SuperLoss",
+    "UncertaintyWeighting",
+    "weighted_loss",
+]

torch_harness-0.10.1/src/torch_harness/layers/__init__.py

@@ -0,0 +1,9 @@
+"""Reliable neural-network layers."""
+
+from torch_harness.layers.safe_batch_norm import (
+    SafeBatchNorm1d,
+    SafeBatchNorm2d,
+    SafeBatchNorm3d,
+)
+
+__all__ = ["SafeBatchNorm1d", "SafeBatchNorm2d", "SafeBatchNorm3d"]

torch_harness-0.10.1/src/torch_harness/layers/safe_batch_norm/DEVLOG.md

@@ -0,0 +1,64 @@
+# SafeBatchNorm development log
+
+## 0.3.0
+
+### Motivating observation
+
+Models often already contain native BatchNorm layers when fault-tolerant
+training is enabled. Runtime conversion should not rely on changing an
+instance's `__class__`, which skips construction and couples correctness to
+Python object-layout details.
+
+### Decision
+
+Add explicit `from_batch_norm` constructors that rebuild each safe layer from
+the corresponding native BatchNorm configuration and persistent state. The
+conversion also preserves training mode and affine-parameter `requires_grad`
+settings. Runtime traversal remains independent through a replacement factory.
+
+### Verified result
+
+Parameterized tests verify that BatchNorm1d, BatchNorm2d, and BatchNorm3d are
+rebuilt as distinct safe instances with equivalent configuration, parameters,
+running statistics, training mode, device, dtype, and `requires_grad` settings.
+The runtime integration also converts nested native BatchNorm layers without
+changing the traversal implementation.
+
+### Unproven claims
+
+The per-layer scalar check introduces a device-to-host synchronization on
+accelerators; its training-throughput cost has not been benchmarked.
+
+## 0.2.0
+
+### Motivating observation
+
+A fault-tolerant training loop can reject non-finite model outputs, but that
+check occurs after BatchNorm has already updated its persistent running state
+from the faulty activation.
+
+### Decision
+
+Directly inherit BatchNorm1d, BatchNorm2d, and BatchNorm3d, and raise
+`FloatingPointError` when a shared pre-forward input check finds a NaN or
+infinity. Checking before BatchNorm prevents any running-state mutation and
+leaves PyTorch's forward and backward implementations untouched for finite
+inputs. Inheritance preserves native constructors, state-dict keys, and type
+identity.
+
+Use an explicit exception rather than `assert`, because optimized Python can
+remove assertions. Synchronized BatchNorm is excluded because a rank-local
+early failure could deadlock peers already entering its collective.
+
+### Verified result
+
+Parameterized tests against BatchNorm1d, BatchNorm2d, and BatchNorm3d verify
+that finite outputs, running-state updates, input gradients, and affine
+parameter gradients match native BatchNorm. NaN and infinity inputs raise
+before the running mean, running variance, or batch counter changes. Evaluation
+mode and disabled running statistics are covered as well.
+
+### Unproven claims
+
+The per-layer scalar check introduces a device-to-host synchronization on
+accelerators; its training-throughput cost has not been benchmarked.