Skip to content

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

Merged
jiahy0825 merged 2 commits intoSandAI-org:mainfrom
wtr0504:refact/joint_partition
Apr 1, 2026
Merged

[refactor]: simplify heuristic save-node selection & add Transformer test#11
jiahy0825 merged 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

@wtr0504 wtr0504 force-pushed the refact/joint_partition branch from 80faf9b to ebea1f0 Compare April 1, 2026 02:42
Copy link
Copy Markdown
Collaborator

@jiahy0825 jiahy0825 left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM

@jiahy0825 jiahy0825 merged commit 719f40e into SandAI-org:main Apr 1, 2026
1 check passed
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.

2 participants