Skip to content

Commit ec68923

Browse files
chen2021673chen2021673
committed
fix: correct distributed rank and unique ID handling
- derive parallel state from the global world size - clarify global rank and per-node process semantics - add multi-node rank regression coverage - restore the NCCL-compatible GetUniqueId interface
1 parent b29714f commit ec68923

12 files changed

Lines changed: 110 additions & 62 deletions

File tree

‎infini_train/include/core/ccl/ccl.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ class CclImpl {
2828

2929
virtual void GetAsyncError(const CclComm *comm, CclStatus *async_error) const;
3030

31-
virtual void CreateUniqueId(CclUniqueId **unique_id, bool generate_id) const;
31+
virtual void GetUniqueId(CclUniqueId **unique_id) const;
3232

3333
virtual void CommInitAll(CclComm **comms, int ndev, const int *devlist) const;
3434

‎infini_train/include/nn/parallel/rank.h‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33
namespace infini_train::nn::parallel {
44
class Rank {
55
public:
6-
Rank(int process_rank, int thread_rank, int process_size, int thread_size);
6+
Rank(int global_process_rank, int thread_rank, int processes_per_node, int threads_per_process);
77

88
int process_rank() const;
99
int thread_rank() const;
@@ -19,9 +19,9 @@ class Rank {
1919
bool IsLastRank() const;
2020

2121
private:
22-
const int process_rank_ = 0; // Rank of the current process within the node
23-
const int thread_rank_ = 0; // Rank of the current thread within the process
24-
const int process_size_ = 1; // Total number of processes on this node
25-
const int thread_size_ = 1; // Total number of threads in the current process
22+
const int global_process_rank_ = 0; // Global rank of the current process
23+
const int thread_rank_ = 0; // Rank of the current thread within the process
24+
const int processes_per_node_ = 1; // Number of processes on each node
25+
const int threads_per_process_ = 1; // Number of threads in each process
2626
};
2727
} // namespace infini_train::nn::parallel

‎infini_train/src/core/ccl/ccl.cc‎

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,7 @@ void CclImpl::GetAsyncError(const CclComm *comm, CclStatus *async_error) const {
1616
LOG(FATAL) << "CclImpl::GetAsyncError is not implemented.";
1717
}
1818

19-
void CclImpl::CreateUniqueId(CclUniqueId **, bool) const {
20-
LOG(FATAL) << "CclImpl::CreateUniqueId is not implemented.";
21-
}
19+
void CclImpl::GetUniqueId(CclUniqueId **unique_id) const { LOG(FATAL) << "CclImpl::GetUniqueId is not implemented."; }
2220

2321
void CclImpl::CommInitAll(CclComm **comms, int ndev, const int *devlist) const {
2422
LOG(FATAL) << "CclImpl::CommInitAll is not implemented.";

‎infini_train/src/core/ccl/cuda/nccl_impl.cc‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -68,16 +68,14 @@ void NcclImpl::GetAsyncError(const CclComm *comm, CclStatus *async_error) const
6868
}
6969
}
7070

71-
void NcclImpl::CreateUniqueId(CclUniqueId **unique_id, bool generate_id) const {
71+
void NcclImpl::GetUniqueId(CclUniqueId **unique_id) const {
7272
CHECK_NOTNULL(unique_id);
7373
if (*unique_id == nullptr) {
7474
*unique_id = new NcclUniqueId();
7575
}
7676
auto *nccl_unique_id = dynamic_cast<NcclUniqueId *>(*unique_id);
7777
CHECK_NOTNULL(nccl_unique_id);
78-
if (generate_id) {
79-
NCCL_CHECK(ncclGetUniqueId(nccl_unique_id->nccl_unique_id()));
80-
}
78+
NCCL_CHECK(ncclGetUniqueId(nccl_unique_id->nccl_unique_id()));
8179
}
8280

8381
void NcclImpl::CommInitAll(CclComm **comms, int ndev, const int *devlist) const {

‎infini_train/src/core/ccl/cuda/nccl_impl.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,7 @@ class NcclImpl final : public CclImpl {
1717

1818
void GetAsyncError(const CclComm *comm, CclStatus *async_error) const override;
1919

20-
void CreateUniqueId(CclUniqueId **unique_id, bool generate_id) const override;
20+
void GetUniqueId(CclUniqueId **unique_id) const override;
2121

2222
void CommInitAll(CclComm **comms, int ndev, const int *devlist) const override;
2323

‎infini_train/src/nn/parallel/global.cc‎

Lines changed: 5 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,6 @@ int GetEnvAsInt(const std::string &name, int default_value) {
1313
return value ? std::atoi(value) : default_value;
1414
}
1515

16-
bool HasEnv(const std::string &name) { return std::getenv(name.c_str()) != nullptr; }
17-
1816
} // namespace
1917

2018
namespace infini_train::nn::parallel::global {
@@ -94,21 +92,15 @@ void GlobalEnv::Init(int nthread_per_process, int tensor_parallel_size, bool seq
9492

9593
CHECK(!initialized_) << "Repeated initialization of GlobalEnv!";
9694

97-
const int proc_world_size = GetEnvAsInt("PROC_WORLD_SIZE", GetEnvAsInt("WORLD_SIZE", 1));
98-
nproc_per_node_ = GetEnvAsInt("NPROC_PER_NODE", GetEnvAsInt("LOCAL_WORLD_SIZE", 1));
95+
const int proc_world_size = GetEnvAsInt("WORLD_SIZE", GetEnvAsInt("PROC_WORLD_SIZE", 1));
96+
nproc_per_node_ = GetEnvAsInt("LOCAL_WORLD_SIZE", GetEnvAsInt("NPROC_PER_NODE", 1));
9997
CHECK_GT(nproc_per_node_, 0) << "NPROC_PER_NODE/LOCAL_WORLD_SIZE must be positive";
10098
CHECK_GT(proc_world_size, 0) << "PROC_WORLD_SIZE/WORLD_SIZE must be positive";
10199
CHECK_EQ(proc_world_size % nproc_per_node_, 0)
102100
<< "PROC_WORLD_SIZE/WORLD_SIZE must be divisible by NPROC_PER_NODE/LOCAL_WORLD_SIZE";
103-
const bool nnodes_env_set = HasEnv("NNODES");
104-
nnodes_ = GetEnvAsInt("NNODES", proc_world_size / nproc_per_node_);
105-
CHECK_GT(nnodes_, 0) << "NNODES must be positive";
106-
if (nnodes_env_set) {
107-
CHECK_EQ(nnodes_ * nproc_per_node_, proc_world_size)
108-
<< "NNODES * NPROC_PER_NODE/LOCAL_WORLD_SIZE must equal PROC_WORLD_SIZE/WORLD_SIZE";
109-
}
110-
global_proc_rank_ = GetEnvAsInt("GLOBAL_PROC_RANK", GetEnvAsInt("RANK", 0));
111-
local_proc_rank_ = GetEnvAsInt("LOCAL_PROC_RANK", GetEnvAsInt("LOCAL_RANK", 0));
101+
nnodes_ = proc_world_size / nproc_per_node_;
102+
global_proc_rank_ = GetEnvAsInt("RANK", GetEnvAsInt("GLOBAL_PROC_RANK", 0));
103+
local_proc_rank_ = GetEnvAsInt("LOCAL_RANK", GetEnvAsInt("LOCAL_PROC_RANK", 0));
112104
CHECK_GE(global_proc_rank_, 0) << "GLOBAL_PROC_RANK/RANK must be non-negative";
113105
CHECK_LT(global_proc_rank_, proc_world_size)
114106
<< "GLOBAL_PROC_RANK/RANK must be less than PROC_WORLD_SIZE/WORLD_SIZE";

‎infini_train/src/nn/parallel/process_group.cc‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -83,13 +83,11 @@ void ProcessGroup::InitMultiProcess(const std::vector<int> &ranks) {
8383
int upper_rank = (global_proc_rank + 1) * n_threads;
8484

8585
core::CclUniqueId *unique_id_raw = nullptr;
86-
87-
int min_rank = std::ranges::min(ranks);
88-
bool is_main_rank = min_rank < upper_rank && min_rank >= lower_rank;
89-
ccl_impl_->CreateUniqueId(&unique_id_raw, is_main_rank);
86+
ccl_impl_->GetUniqueId(&unique_id_raw);
9087
std::unique_ptr<core::CclUniqueId> unique_id(unique_id_raw);
9188

92-
if (is_main_rank) {
89+
int min_rank = std::ranges::min(ranks);
90+
if (min_rank < upper_rank && min_rank >= lower_rank) {
9391
is_main_process_ = true;
9492
core::WriteUniqueIdFile(*unique_id, name_);
9593
} else {

‎infini_train/src/nn/parallel/rank.cc‎

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -2,18 +2,19 @@
22
#include "infini_train/include/nn/parallel/global.h"
33

44
namespace infini_train::nn::parallel {
5-
Rank::Rank(int process_rank, int thread_rank, int process_size, int thread_size)
6-
: process_rank_(process_rank), thread_rank_(thread_rank), process_size_(process_size), thread_size_(thread_size) {}
5+
Rank::Rank(int global_process_rank, int thread_rank, int processes_per_node, int threads_per_process)
6+
: global_process_rank_(global_process_rank), thread_rank_(thread_rank), processes_per_node_(processes_per_node),
7+
threads_per_process_(threads_per_process) {}
78

8-
int Rank::process_rank() const { return process_rank_; }
9+
int Rank::process_rank() const { return global_process_rank_; }
910
int Rank::thread_rank() const { return thread_rank_; }
1011

11-
int Rank::process_size() const { return process_size_; }
12-
int Rank::thread_size() const { return thread_size_; }
12+
int Rank::process_size() const { return processes_per_node_; }
13+
int Rank::thread_size() const { return threads_per_process_; }
1314

14-
int Rank::GlobalRank() const { return process_rank_ * thread_size_ + thread_rank_; }
15+
int Rank::GlobalRank() const { return global_process_rank_ * threads_per_process_ + thread_rank_; }
1516

16-
bool Rank::IsParallel() const { return thread_size_ * process_size_ > 1; }
17+
bool Rank::IsParallel() const { return global::GetWorldSize() > 1; }
1718

1819
bool Rank::IsMainRank() const { return GlobalRank() == 0; }
1920

‎tests/CMakeLists.txt‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,9 @@ include(${CMAKE_SOURCE_DIR}/cmake/test_macros.cmake)
77
# Common test utilities
88
add_subdirectory(common)
99

10+
# Distributed tests
11+
add_subdirectory(distributed)
12+
1013
# DataLoader tests
1114
add_subdirectory(dataloader)
1215

‎tests/distributed/CMakeLists.txt‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
# ==========================================================================
2+
# Distributed tests
3+
# ==========================================================================
4+
5+
infini_train_add_test(test_rank
6+
SOURCES test_rank.cc
7+
LABELS cpu
8+
)
9+
10+
add_test(
11+
NAME RankTest.MultiNodeSingleProcessIsParallel
12+
COMMAND ${CMAKE_COMMAND} -E env
13+
"WORLD_SIZE=2"
14+
"LOCAL_WORLD_SIZE=1"
15+
"RANK=0"
16+
"LOCAL_RANK=0"
17+
$<TARGET_FILE:test_rank>
18+
--gtest_filter=RankTest.DetectsParallelismFromGlobalWorldSize
19+
)
20+
21+
set_tests_properties(RankTest.MultiNodeSingleProcessIsParallel
22+
PROPERTIES
23+
LABELS cpu
24+
TIMEOUT 10
25+
)

0 commit comments

Comments
 (0)