dist_config.hpp 419 Bytes
Newer Older
1
2
#pragma once

3
#include <string>
4
5
6
7
8
9
10
11
12
13
14
15
#include <vector>

namespace infinilm::engine::distributed {

struct DistConfig {
    // Device IDs for each rank in tensor parallelism
    std::vector<int> tp_device_ids;

    DistConfig();
    explicit DistConfig(int tp_size);
    explicit DistConfig(const std::vector<int> &tp_device_ids_);

16
    explicit operator std::string() const;
17
18
19
};

} // namespace infinilm::engine::distributed