Fix/training error publish - #2
Open
SLCZ687 wants to merge 2 commits into
Open
Conversation
There was a problem hiding this comment.
Pull request overview
This PR focuses on making Metis training and checkpointing more reliable under DeepSpeed ZeRO-3—especially for “delta” checkpoints—while also improving memory behavior and tightening release/audit hygiene.
Changes:
- Add ZeRO-3-aware delta save/load paths that avoid full-model consolidation and support sharded parameter initialization/copying.
- Reorder training setup to construct
TrainingArgumentsbefore model load (for ZeRO-3 partition-on-init), add--save_total_limit, and adjust step-0 checkpoint behavior under DeepSpeed. - Improve operational ergonomics: mmap-backed dataset cache loads, safer code snapshotting, release text scanning updates, and environment/README tweaks.
Reviewed changes
Copilot reviewed 15 out of 15 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| train/trainer.py | Overrides save_model for ZeRO-3 + delta to avoid full consolidation; adds a guard rail in _save. |
| train/train_utils.py | Makes parameter counting ZeRO-3-aware; hardens code snapshot filtering to avoid copying outputs. |
| train/run_train.py | Builds TrainingArguments before model construction; adds save_total_limit, ZeRO-3 load logging, and skips step-0 checkpoint under DeepSpeed. |
| train/dataset.py | Uses torch.load(..., mmap=True) with fallbacks to reduce host RAM duplication across ranks. |
| scripts/train.sh | Adds --save-total-limit option plumbing and validation. |
| scripts/run_train.sh | Passes --save_total_limit through to the Python training entrypoint. |
| README.md | Documents post-conda pip installs for CUDA extensions. |
| README_zh.md | Mirrors README CUDA-extension install guidance in Chinese. |
| metis/weight_utils.py | Adds ZeRO-3 detection and shard-safe state copy/initialization helpers. |
| metis/checkpoint_utils.py | Implements per-parameter ZeRO-3 delta gathering and ZeRO-3-aware delta loading. |
| metis/backbone_wrappers/Qwen3_5_wrapper.py | Switches to a ZeRO-3-safe (reentrant) checkpointing wrapper under grad-enabled paths. |
| eval/tests/test_protocols.py | Tightens release text scanning for private paths and dev repo names; expands suffix coverage. |
| eval/experiments/ablation/run.py | Removes hardcoded /tmp literal and renames dev-root metadata key to more generic model repo naming. |
| eval/configs/assets.json | Updates public note to avoid development repo naming. |
| environment.yml | Adjusts CUDA package sourcing/pinning and documents CUDA-extension build ordering. |
Suppressed comments (1)
metis/checkpoint_utils.py:172
- The post-save
torch.distributed.barrier()is unguarded. This can crash in non-distributed contexts (or before process-group init) even though saving is otherwise local I/O.
if uses_zero3_collectives:
torch.distributed.barrier()
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+133
to
+136
| if not is_main_process: | ||
| if uses_zero3_collectives: | ||
| torch.distributed.barrier() | ||
| return |
Comment on lines
+67
to
+70
| import deepspeed | ||
|
|
||
| with deepspeed.zero.GatheredParameters(zero_parameters, modifier_rank=0): | ||
| yield torch.distributed.get_rank() == 0 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
No description provided.