Skip to content

Incomplete Parameter Collection During Optimization Passes #2

@eric-tramel

Description

@eric-tramel

Description

During optimization passes, not all Parameters are being updated. This suggests that the parameter collection logic is incomplete - some Parameters that should receive updates are being missed during the backward pass or optimizer step.

Potential root causes:

  1. named_parameters() or parameters() methods not collecting all nested Parameters
  2. Parameters in certain module structures (nested modules, lists, dicts) being skipped
  3. Graph traversal during backward pass missing certain Parameter nodes
  4. Feedback not being propagated to all Parameters that require gradients

All Parameters with requires_grad=True that contributed to the forward pass should receive feedback and be eligible for updates during optimization.

Acceptance Criteria

  • Investigate named_parameters() and parameters() in InferenceModule for collection completeness
  • Verify that nested module Parameters are correctly discovered
  • Verify that Parameters in container structures (lists, dicts) are collected
  • Ensure backward pass visits all Parameter nodes in the computation graph
  • Ensure feedback accumulation reaches all relevant Parameters
  • Add tests for nested module Parameter collection
  • Add tests for container-based Parameter collection
  • Add integration test verifying all Parameters are updated after optimization
  • Verify fix with cookbooks/hallucination_detection.py

Notes

  • Key files: module.py (InferenceModule), parameter.py, optimization/backward.py, optimization/optimizer.py
  • The torch.nn.Module pattern for parameter discovery should be followed closely
  • Consider adding debugging/logging to track which Parameters are collected vs. updated

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions