mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-06-02 20:35:29 +00:00
554aa9592e
* [legacy] move communication to legacy (#4640) * [legacy] refactor logger and clean up legacy codes (#4654) * [legacy] make logger independent to gpc * [legacy] make optim independent to registry * [legacy] move test engine to legacy * [legacy] move nn to legacy (#4656) * [legacy] move nn to legacy * [checkpointio] fix save hf config * [test] remove useledd rpc pp test * [legacy] fix nn init * [example] skip tutorial hybriad parallel example * [devops] test doc check * [devops] test doc check
16 lines
483 B
Python
16 lines
483 B
Python
#!/usr/bin/env python
|
|
# -*- encoding: utf-8 -*-
|
|
|
|
|
|
def _calc_incoming_device_range(i, rank, world_size, sub_seq_length):
|
|
device_of_incoming_k = (rank - i - 1) % world_size
|
|
start_idx = sub_seq_length * device_of_incoming_k
|
|
end_idx = sub_seq_length * (device_of_incoming_k + 1)
|
|
return start_idx, end_idx
|
|
|
|
|
|
def _calc_current_device_range(rank, sub_seq_length):
|
|
start_idx = sub_seq_length * rank
|
|
end_idx = sub_seq_length * (rank + 1)
|
|
return start_idx, end_idx
|