Conversation
Signed-off-by: 0z5a <192209249+0z5a@users.noreply.github.com>
![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) |
1 participant
Problem
Native Torch Adam/AdamW checkpoint resume fails at several boundaries:
step, although the scalar is restored from common checkpoint state.Changes
Update the Megatron Docker patch to accept the empty template, restore independent per-parameter step tensors, remove the checkpoint-only group field and retain the restored step while copying parameter shards.
Add regressions using the real Megatron methods and a fixed-batch comparison runner: eight uninterrupted updates versus four updates followed by four updates in a fresh process, including training offload and wake-up.
Validation
1dcf0dafa884ad52ffb243625717a3471643e087; resulting optimizer source matches the tested runtimeThe fixture uses full Qwen3-0.6B with 596,049,920 parameters, BF16 model parameters, FP32 master parameters, native distributed AdamW and TP=PP=CP=1. The runner checks parameters, moments, steps, scheduler, RNG, gradients, metrics and saved training data. Missing moments and inconsistent steps are negative controls.
Draft status and limitations
No performance, convergence, larger-model, Transformer Engine or topology-resharding improvement is claimed.