gridfm / gridfm/gridfm-graphkit
LayeredWeightedPhysicsLoss reads layer_residuals off the model instead of receiving it
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 105
- Forks
- 36
- Avg merge
- 1d 8h
- Merged PRs (30d)
- 9
Description
Low priority / cleanup. Pre-existing design issue, no known incorrect results today. Surfaced while reviewing #87.
GNS_heterogeneous.forward writes self.layer_residuals as a side effect, and LayeredWeightedPhysicsLoss (training/loss.py:363) reads model.layer_residuals afterwards. The data dependency is invisible in both signatures — the loss just receives model=self.model.
Concrete downsides:
- Unenforced ordering. Requires
forward()beforeloss_fn()on the same instance. If the dict is empty,L = 0produces no loop iterations,total_lossstays a Python float, and it fails ontotal_loss.item()— an obscure error far from the cause. - Not reuse-safe. One dict per module instance: two forward passes before a loss call silently overwrite, so gradient accumulation or evaluating two batches before reducing would read the wrong batch's residuals and return a plausible number rather than erroring.
- Hidden model/loss coupling. Only works with models exposing
layer_residuals, i.e.GNS_heterogeneous;AttributeErroronGRIT. A config pairing GRIT with this loss fails at runtime, not at validation.
Working as intended today: the stored tensors keep their autograd graph (no .detach()), so gradients flow correctly.
Suggested fix: return residuals from forward alongside the predictions (as #87 now does for embeddings) and pass them into the loss explicitly, so the dependency is in the signature and the model= kwarg is no longer needed for this.
Deferred from #87 because the fix touches the shared BaseLoss.forward signature across all losses.
Refs #87
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start with GNS_heterogeneous.forward and LayeredWeightedPhysicsLoss in training/loss.py:363, then inspect the shared BaseLoss.forward signature and the changes referenced in #87. Trace how residuals are produced and consumed, including other losses affected by the shared signature. Done means residuals are passed explicitly, the model= dependency is removed, and the loss no longer relies on stored per-instance state.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100