Skip to content

[refactor]: simplify heuristic save-node selection & add Transformer test#11

Open
wtr0504 wants to merge 2 commits intoSandAI-org:mainfrom
wtr0504:refact/joint_partition
Open

[refactor]: simplify heuristic save-node selection & add Transformer test#11
wtr0504 wants to merge 2 commits intoSandAI-org:mainfrom
wtr0504:refact/joint_partition

Conversation

@wtr0504
Copy link
Copy Markdown
Collaborator

@wtr0504 wtr0504 commented Mar 31, 2026

🗂️ PR Category

  • ✨ New Feature
  • 🚀 Optimization (performance, memory, etc.)
  • 💥 Breaking Change
  • 🐛 Bug Fix
  • 🛠️ Development / Refactoring
  • 📚 Documentation
  • 🧹 Chore (Dependencies, CI/CD, Configuration, etc.)
  • 🧪 Testing

📝 Description

  • Unify primal and compute-intensive node handling into a single
    _decide_save_node pass over the joint graph, replacing two
    separate loops
  • Fix infinite-loop bug in _primal_contributes_to_bwd_directly by
    adding a visited set to the BFS worklist
  • Remove global SAVE_TENSOR_NODES; track saved nodes locally via
    closure in custom_joint_graph_partition_fn
  • Deduplicate with ctx: block for HANDCRAFT and HEURISTIC policies
  • Add TransformerConfig, Transformer, and related helpers to
    model_definition.py for testing

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant