[Bug]: GPU OOM when capturing graph.

[TianyiZhao1437]

Steps to Reproduce

1. Use 4070 12GB.
2. Join parallax gpt-oss-20b.
3. Will be assigned with 3 layers of ~8 GB weight usage.
4. OOM when capturing graph. Only ~0.5GB for graph.

Expected Behavior

Should be assigned with less layers. Finish capturing graph which takes around 2GB memory.

Actual Behavior

OOM when capturing graph.

Version

v0.0.1

Environment & Context

• I'm using the latest version.
• I have searched existing issues.