Skip to content

Add support of tp_size != WORLD_SIZE to bloom #428

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
regisss merged 2 commits into from
Oct 23, 2023
Merged

Add support of tp_size != WORLD_SIZE to bloom #428

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f5690e0
Add support of tp_size != WORLD_SIZE to bloom
NirSonnenschein Sep 27, 2023
636f1a8
Update optimum/habana/transformers/models/bloom/modeling_bloom.py
NirSonnenschein Oct 23, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 16 additions & 3 deletions optimum/habana/transformers/models/bloom/modeling_bloom.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,8 +85,13 @@ def gaudi_bloom_build_alibi_tensor(
).unsqueeze(0).expand(num_heads, -1, -1)

# Select the part of the tensor that corresponds to our tensor parallel index.
tp_world_size = int(os.environ.get("WORLD_SIZE", 1))
tp_index = int(os.environ.get("RANK", 0))
# if inference_tp_size is set use it instead of world size
world = int(os.environ.get("WORLD_SIZE", 1))
tp_world_size = GaudiBloomForCausalLM.inference_tp_size if GaudiBloomForCausalLM.inference_tp_size else world
tp_index = 0 # if world size == 1 ignore rank and use 0 (for cases where WORLD_SIZE is not equal to tp size)
if tp_world_size > 1:
tp_index = int(os.environ.get("RANK", 0))

alibi = alibi.reshape((tp_world_size, -1, *alibi.shape[1:]))[tp_index]

alibi = alibi.repeat(batch_size, 1, 1)
Expand Down Expand Up @@ -282,7 +287,8 @@ def gaudi_bloom_convert_to_standard_cache(
if training:
num_heads = batch_size_times_num_heads // batch_size
else:
tp_world_size = int(os.environ.get("WORLD_SIZE", 1))
world = int(os.environ.get("WORLD_SIZE", 1))
tp_world_size = GaudiBloomForCausalLM.inference_tp_size if GaudiBloomForCausalLM.inference_tp_size else world
num_heads = self.config.n_head // tp_world_size
batch_size = batch_size_times_num_heads // num_heads
# key: [batch_size * num_heads, head_dim, seq_length] -> [batch_size, num_heads, head_dim, seq_length]
Expand Down Expand Up @@ -462,6 +468,13 @@ def custom_forward(*inputs):


class GaudiBloomForCausalLM(BloomForCausalLM):
inference_tp_size = None
Copy link
Copy Markdown
Collaborator

@regisss regisss Sep 27, 2023

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would rather use a TP_SIZE env variable than adding a class attribute here


def set_tp_for_inference(tp_for_inference: int):
world = int(os.environ.get("WORLD_SIZE", 1))
assert tp_for_inference == 1 or tp_for_inference == world, "only setting 1 (no tp) or world size is supported"
GaudiBloomForCausalLM.inference_tp_size = tp_for_inference

def prepare_inputs_for_generation(
self,
input_ids: torch.LongTensor,
Expand Down