mirror of https://github.com/hpcaitech/ColossalAI
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
|