Repository navigation
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request updates the Megatron patch to properly handle fresh native Torch Adam/AdamW optimizers during checkpoint loading, ensuring that steps are independent and updated correctly. It also introduces a comprehensive suite of integration and unit tests to verify checkpoint resume parity. The review feedback highlights a few improvement opportunities: cleaning up the temporary directory ray_temp in the parity runner to prevent disk space accumulation, and using weights_only=True when loading PyTorch checkpoints in the audit tests to enhance security and avoid deprecation warnings.
| try: | ||
| ray_temp = tempfile.mkdtemp(prefix="vime-resume-ray-") | ||
| (output / "ray-temp-dir.txt").write_text(ray_temp + "\n") | ||
| ray.init( | ||
| address="local", | ||
| num_cpus=12, | ||
| num_gpus=1, | ||
| include_dashboard=False, | ||
| object_store_memory=2 * 1024**3, | ||
| namespace=output.name, | ||
| _temp_dir=ray_temp, | ||
| _node_ip_address="127.0.0.1", | ||
| runtime_env={"env_vars": {"VIME_RESUME_AUDIT_DIR": str(output)}}, | ||
| ) | ||
| spec = importlib.util.spec_from_file_location("resume_parity_train", checkout / "train.py") | ||
| module = importlib.util.module_from_spec(spec) | ||
| spec.loader.exec_module(module) | ||
| train_args = module.parse_args() | ||
| if args.stop_after is not None: | ||
| assert 0 < args.stop_after < train_args.num_rollout | ||
| # Do not change num_rollout: the scheduler horizon must be identical | ||
| # in the continuous run, split prefix, and resumed suffix. | ||
| module.range = lambda start, stop: range(start, min(stop, args.stop_after)) | ||
| module.train(train_args) | ||
| (output / "train-returned.json").write_text(json.dumps({"completed": True})) | ||
| finally: | ||
| ray.shutdown() | ||
| sys.stdout.flush() | ||
| sys.stderr.flush() | ||
| for fd, original in zip([1, 2], saved_fds, strict=True): | ||
| os.dup2(original, fd) | ||
| os.close(original) | ||
| log.close() |
There was a problem hiding this comment.
The temporary directory ray_temp created by tempfile.mkdtemp is not cleaned up after the training and audit logs are recorded. This can lead to disk space accumulation in /tmp on test machines or CI runners over time.
Please clean up the temporary directory in the finally block using shutil.rmtree.
ray_temp = tempfile.mkdtemp(prefix="vime-resume-ray-")
try:
(output / "ray-temp-dir.txt").write_text(ray_temp + "\n")
ray.init(
address="local",
num_cpus=12,
num_gpus=1,
include_dashboard=False,
object_store_memory=2 * 1024**3,
namespace=output.name,
_temp_dir=ray_temp,
_node_ip_address="127.0.0.1",
runtime_env={"env_vars": {"VIME_RESUME_AUDIT_DIR": str(output)}},
)
spec = importlib.util.spec_from_file_location("resume_parity_train", checkout / "train.py")
module = importlib.util.module_from_spec(spec)
spec.loader.exec_module(module)
train_args = module.parse_args()
if args.stop_after is not None:
assert 0 < args.stop_after < train_args.num_rollout
# Do not change num_rollout: the scheduler horizon must be identical
# in the continuous run, split prefix, and resumed suffix.
module.range = lambda start, stop: range(start, min(stop, args.stop_after))
module.train(train_args)
(output / "train-returned.json").write_text(json.dumps({"completed": True}))
finally:
ray.shutdown()
import shutil
shutil.rmtree(ray_temp, ignore_errors=True)
sys.stdout.flush()
sys.stderr.flush()
for fd, original in zip([1, 2], saved_fds, strict=True):
os.dup2(original, fd)
os.close(original)
log.close()| a = torch.load(continuous / f"dumps/train_data/{rid}.pt", map_location="cpu", weights_only=False) | ||
| b = torch.load(root / f"dumps/train_data/{rid}.pt", map_location="cpu", weights_only=False) |
There was a problem hiding this comment.
Using weights_only=False when loading PyTorch checkpoints is a security risk and can trigger deprecation warnings in newer PyTorch versions. Since the saved training data only contains standard PyTorch tensors and basic Python types, it is safer and cleaner to use weights_only=True.
| a = torch.load(continuous / f"dumps/train_data/{rid}.pt", map_location="cpu", weights_only=False) | |
| b = torch.load(root / f"dumps/train_data/{rid}.pt", map_location="cpu", weights_only=False) | |
| a = torch.load(continuous / f"dumps/train_data/{rid}.pt", map_location="cpu", weights_only=True) | |
| b = torch.load(root / f"dumps/train_data/{rid}.pt", map_location="cpu", weights_only=True) |
| @pytest.mark.parametrize("key", ["tokens", "loss_masks", "log_probs", "advantages"]) | ||
| def test_rejects_changed_training_batch(evidence, key): | ||
| path = evidence[2] / "dumps/train_data/1.pt" | ||
| batch = torch.load(path, weights_only=False) |
There was a problem hiding this comment.
Using weights_only=False when loading PyTorch checkpoints is a security risk and can trigger deprecation warnings in newer PyTorch versions. Since the saved training data only contains standard PyTorch tensors and basic Python types, it is safer and cleaner to use weights_only=True.
| batch = torch.load(path, weights_only=False) | |
| batch = torch.load(path, weights_only=True) |
d7ea0b7 to
a3098e0
Compare
Signed-off-by: 0z5a <192209249+0z5a@users.noreply.github.com> Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
a3098e0 to
f8ee091
Compare
Signed-off-by: 0z5a <dezhen.lu@student.uni-tuebingen.de>
|
Superseded by #460, now merged into |
Incremental review: Files changed.
Base:
0z5a-branchat791f101e20db377b28e15f12a00863baaea2605f.Fix native Torch Adam/AdamW checkpoint restoration in the existing Megatron Docker patch. Accept an empty fresh-optimizer load template, skip shard-copying the common scalar step, restore independent per-parameter step tensors, and remove the checkpoint-only group field.
The cleaned branch incorporates the current target base. The submitted diff contains only the Docker patch; audit runners and added test files were removed.
1dcf0dafa884ad52ffb243625717a3471643e087The fixture uses full Qwen3-0.6B, BF16 model parameters, FP32 master parameters, distributed AdamW and TP=PP=CP=1. The final aggregate comparison remains unverified; additional online checkpoint-resume validation remains incomplete. The GPU fixture includes #426 as a separate runtime prerequisite. The cleanup preserves the implementation while merging current base changes; the archived GPU results precede this cleanup.
No speed or convergence improvement is established.