fix actor init bugs #1098
fix actor init bugs #1098
Conversation
There was a problem hiding this comment.
Pull request overview
This PR moves the torch_memory_saver initialization logic from the generic TrainRayActor initialization to the Megatron-specific initialization function, addressing an initialization timing issue.
Key changes:
- Removed
torch_memory_saverimport and configuration fromslime/ray/train_actor.py - Added
torch_memory_saverimport and configuration toslime/backends/megatron_utils/initialize.py
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
slime/ray/train_actor.py |
Removed torch_memory_saver import and memory margin configuration logic from TrainRayActor.init() |
slime/backends/megatron_utils/initialize.py |
Added torch_memory_saver import and memory margin configuration to Megatron initialization function |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| if args.offload_train: | ||
| if (x := args.train_memory_margin_bytes) > 0: | ||
| logger.info(f"Set torch_memory_saver.memory_margin_bytes to {x}") | ||
| torch_memory_saver.memory_margin_bytes = x |
There was a problem hiding this comment.
The logic has been inverted compared to the original implementation. In the original code in train_actor.py, if train_memory_margin_bytes was greater than 0, it would assert that offload_train must be true, raising an error if it wasn't. The new implementation silently ignores train_memory_margin_bytes when offload_train is false. Consider whether this silent failure is intentional or if validation should still occur when train_memory_margin_bytes is set without offload_train enabled.
2 participants
No description provided.