merge trace file (#31)

wenxie-amd · web-flow · commit b12ad8851045 · 2025-04-21T13:17:23.000+08:00
diff --git a/examples/megatron/run_pretrain.sh b/examples/megatron/run_pretrain.sh
@@ -189,7 +189,7 @@ elif [ "$PRIMUS_HIPBLASLT_TUNING_STAGE" -eq 2 ]; then
             exit 1
         fi
 
-        bash "${PRIMUS_PATH}"/examples/scripts/docker_podman_proxy.sh run --rm \
+        bash "${PRIMUS_PATH}"/tools/docker/docker_podman_proxy.sh run --rm \
             --env HIP_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES \
             --env GPU_MAX_HW_QUEUES=$GPU_MAX_HW_QUEUES \
             --env TORCH_NCCL_HIGH_PRIORITY=$TORCH_NCCL_HIGH_PRIORITY \
@@ -299,7 +299,7 @@ elif [ "$RUN_ENV" = "slurm" ]; then
         "
     fi
 
-    bash "${PRIMUS_PATH}"/examples/scripts/docker_podman_proxy.sh run --rm \
+    bash "${PRIMUS_PATH}"/tools/docker/docker_podman_proxy.sh run --rm \
         --env SLURM_MASTER_ADDR=$SLURM_MASTER_ADDR \
         --env SLURM_MASTER_PORT=$SLURM_MASTER_PORT \
         --env SLURM_PROCID=$SLURM_PROCID \
diff --git a/examples/megatron/run_slurm_pretrain.sh b/examples/megatron/run_slurm_pretrain.sh
@@ -9,8 +9,8 @@
 export RUN_ENV=slurm
 export MODEL_CONFIG=deepseek_v2_lite
 
-export PRIMUS_HIPBLASLT_TUNING_STAGE=${PRIMUS_HIPBLASLT_TUNING_STAGE:-3}
-export NUM_NODES=${NUM_NODES:-8}
+export PRIMUS_HIPBLASLT_TUNING_STAGE=${PRIMUS_HIPBLASLT_TUNING_STAGE:-0}
+export NUM_NODES=${NUM_NODES:-1}
 
 SCRIPT_DIR=$(dirname "$(realpath "${BASH_SOURCE[0]}")")
 
diff --git a/tools/docker/docker_podman_proxy.sh b/tools/docker/docker_podman_proxy.sh
diff --git a/tools/docker/start_container.sh b/tools/docker/start_container.sh
@@ -0,0 +1,24 @@
+#!/bin/bash
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+#################################################################################
+
+PRIMUS_PATH=$(realpath "$(dirname "$0")/../..")
+export DOCKER_IMAGE="docker.io/rocm/megatron-lm:latest"
+
+bash "${PRIMUS_PATH}"/tools/docker/docker_podman_proxy.sh run -d \
+    --name dev_primus \
+    --ipc=host \
+    --network=host \
+    --device=/dev/kfd \
+    --device=/dev/dri \
+    --device=/dev/infiniband \
+    --cap-add=SYS_PTRACE \
+    --cap-add=CAP_SYS_ADMIN \
+    --security-opt seccomp=unconfined \
+    --group-add video \
+    --privileged \
+    -v "${PRIMUS_PATH}:${PRIMUS_PATH}" \
+    $DOCKER_IMAGE sleep infinity
diff --git a/tools/preflight/global_vars.py b/tools/preflight/global_vars.py
@@ -1,3 +1,9 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+#################################################################################
+
 import os
 
 WORLD_SIZE = int(os.environ.get("WORLD_SIZE", 1))
diff --git a/tools/preflight/inter_node_comm.py b/tools/preflight/inter_node_comm.py
@@ -1,3 +1,9 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+#################################################################################
+
 import time
 
 import matplotlib.pyplot as plt
diff --git a/tools/preflight/inter_node_comm_p2p.py b/tools/preflight/inter_node_comm_p2p.py
@@ -1,3 +1,9 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+#################################################################################
+
 import time
 
 import matplotlib.pyplot as plt
diff --git a/tools/preflight/intra_node_comm.py b/tools/preflight/intra_node_comm.py
@@ -1,3 +1,9 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+#################################################################################
+
 import time
 
 import matplotlib.pyplot as plt
diff --git a/tools/preflight/preflight_perf_test.py b/tools/preflight/preflight_perf_test.py
@@ -1,3 +1,9 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+#################################################################################
+
 import argparse
 import os
 
diff --git a/tools/preflight/square_gemm.py b/tools/preflight/square_gemm.py
@@ -1,3 +1,9 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+#################################################################################
+
 import time
 
 import matplotlib.pyplot as plt
diff --git a/tools/preflight/utility.py b/tools/preflight/utility.py
@@ -1,3 +1,9 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+#################################################################################
+
 import os
 import socket
 from pathlib import Path
diff --git a/tools/profile_trace/merge_trace_file.py b/tools/profile_trace/merge_trace_file.py
@@ -0,0 +1,39 @@
+###############################################################################
+# Copyright (c) 2025, Advanced Micro Devices, Inc. All rights reserved.
+#
+# See LICENSE for license information.
+#################################################################################
+
+import argparse
+import glob
+import json
+import os
+
+
+def merge_traces_in_folder(path, output_file):
+    merged_trace = {"traceEvents": [], "displayTimeUnit": "ns"}
+
+    json_files = sorted(glob.glob(os.path.join(path, "*.json")))
+
+    if not json_files:
+        print("No JSON files found in the folder.", flush=True)
+        return
+
+    for filepath in json_files:
+        with open(filepath, "r") as f:
+            trace = json.load(f)
+            merged_trace["traceEvents"].extend(trace.get("traceEvents", []))
+
+    with open(output_file, "w") as f:
+        json.dump(merged_trace, f)
+
+    print(f"Merged trace saved to: {output_file}", flush=True)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--profile-trace-path", type=str)
+    parser.add_argument("--merged-trace-file", type=str)
+    args = parser.parse_args()
+
+    merge_traces_in_folder(args.profile_trace_path, args.merged_trace_file)
