add nemo run export

Signed-off-by: Jennifer Chen <[email protected]>

Author: jenchen13

examples/nemo_run/common/process_openscience.py

@@ -45,49 +45,12 @@ def process_subset(raw_dir, proc_dir):
     split_ds["test"].to_json(os.path.join(proc_dir, "validation.jsonl"))
 
 
-# TODO remove below?
-def sample_openscience(raw_dir, proc_dir, sample_ratio=1):
-    """Process raw OpenScience data by subsampling the dataset by default, then
-    writing into train/val split with 99/1 ratio"""
-    files = os.listdir(raw_dir)
-    num_data = 0
-
-    for file in files:
-        # Open each jsonl
-        if file.endswith("jsonl"):
-            print(f"Sampling from {file}")
-            with (
-                open(os.path.join(raw_dir, file)) as f_raw,
-                open(os.path.join(proc_dir, "training.jsonl"), "a") as f_train,
-                open(os.path.join(proc_dir, "validation.jsonl"), "a") as f_val,
-            ):
-                for idx, line in enumerate(f_raw):
-                    if idx % sample_ratio != 0:
-                        continue
-                    data = json.loads(line)
-                    # convert dictionary to OpenAI chat: from {"input": "...", "output": "..."}
-                    # to [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]
-                    data = {
-                        "messages": [
-                            {"role": "user", "content": data["input"]},
-                            {"role": "assistant", "content": data["output"]},
-                        ]
-                    }
-
-                    if num_data % 100 == 0:
-                        f_val.write(json.dumps(data) + "\n")
-                    else:
-                        f_train.write(json.dumps(data) + "\n")
-                    num_data += 1
-
-
 if __name__ == "__main__":
     args = get_parser().parse_args()
     raw_dir = f"{args.output_dir}/openscience_raw"
     proc_dir = f"{args.output_dir}/openscience_proc"
 
     if not os.path.exists(raw_dir):
-        # download_hf_dataset("nvidia/OpenScience", raw_dir)
         q235_subset = load_dataset("nvidia/OpenScience", data_files="OS-Q3-235B-4.jsonl")
         q235_subset.save_to_disk(raw_dir)
 
@@ -97,4 +60,3 @@ def sample_openscience(raw_dir, proc_dir, sample_ratio=1):
         process_subset(raw_dir, proc_dir)
     else:
         print(f"Processed OpenScience dataset exists in: {proc_dir}, skipped processing")
-    # process_openscience(raw_dir, proc_dir)

modelopt/torch/export/plugins/nemo_run.py

@@ -0,0 +1,71 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Export functions for NeMo Run."""
+
+from pathlib import Path
+
+from nemo.collections.llm.api import export_ckpt
+from nemo.utils import logging
+
+
+def export_most_recent_ckpt(directory: str, output_path: str):
+    """Export most recent checkpoint from a NeMo Run experiment directory."""
+    most_recent_ckpt = _get_most_recent_ckpt(directory)
+    modelopt_kwargs = {"export_extra_modules": True}
+    logging.info(f"Exporting most recent NeMo Run checkpoint: {most_recent_ckpt}")
+    export_ckpt(
+        most_recent_ckpt,
+        "hf",
+        output_path=output_path,
+        overwrite=True,
+        modelopt_export_kwargs=modelopt_kwargs,
+    )
+
+
+def _get_most_recent_subdir(directory: Path):
+    # Get all subdirectories
+    subdirs = [d for d in directory.iterdir() if d.is_dir()]
+    if not subdirs:
+        raise ValueError(f"No subdirectories found in {directory}")
+
+    # Sort by modification time (most recent first)
+    most_recent = max(subdirs, key=lambda x: x.stat().st_mtime)
+
+    return most_recent
+
+
+def _get_most_recent_ckpt(directory: str):
+    """Find the most recent checkpoint subdirectory in a given NeMo Run experiment directory.
+
+    Args:
+        directory (str): Path to the directory to search in.
+
+    Returns:
+        str: Path to the most recent subdirectory.
+    """
+    exp_dir = Path(directory) / "default"
+    assert exp_dir.exists(), f"Experiment directory {exp_dir} does not exist"
+
+    checkpoint_dir = exp_dir / "checkpoints"
+    if checkpoint_dir.exists():
+        most_recent = _get_most_recent_subdir(checkpoint_dir)
+    else:
+        most_recent = _get_most_recent_subdir(exp_dir)
+        checkpoint_dir = most_recent / "checkpoints"
+        assert checkpoint_dir.exists(), f"Checkpoint directory {checkpoint_dir} does not exist"
+        most_recent = _get_most_recent_subdir(checkpoint_dir)
+
+    return str(most_recent)