Skip to content

Commit 8dee68a

Browse files
CjhHa1CjhHa1
committed
bloom policy
1 parent 0192011 commit 8dee68a

1 file changed

Lines changed: 219 additions & 1 deletion

File tree

‎colossalai/shardformer/policies/bloom.py‎

Lines changed: 219 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,26 @@
1+
import warnings
2+
from functools import partial
3+
from types import MethodType
4+
from typing import Dict, List, Optional, Tuple, Union
5+
6+
import numpy as np
7+
import torch
18
import torch.nn as nn
9+
from torch import Tensor
10+
from torch.nn import CrossEntropyLoss, Module
11+
from transformers.modeling_outputs import BaseModelOutputWithPastAndCrossAttentions
12+
from transformers.models.bloom.modeling_bloom import BloomModel
13+
from transformers.utils import logging
214

315
import colossalai.shardformer.layer as col_nn
16+
from colossalai.pipeline.stage_manager import PipelineStageManager
417

518
from .._utils import getattr_, setattr_
619
from ..modeling.bloom import build_bloom_alibi_tensor_fn
720
from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
821

22+
logger = logging.get_logger(__name__)
23+
924

1025
class BloomPolicy(Policy):
1126

@@ -110,7 +125,40 @@ def postprocess(self):
110125

111126

112127
class BloomModelPolicy(BloomPolicy):
113-
pass
128+
129+
def __init__(self) -> None:
130+
super().__init__()
131+
132+
def module_policy(self):
133+
module_policy = super().module_policy()
134+
from transformers.models.bloom.modeling_bloom import BloomModel
135+
if self.pipeline_stage_manager:
136+
module_policy[BloomModel] = ModulePolicyDescription(
137+
method_replacement={"forward": partial(bloom_model_forward, stage_manager=self.pipeline_stage_manager)})
138+
139+
def get_held_layers(self) -> List[Module]:
140+
"""
141+
get pipeline layers for current stage
142+
"""
143+
module = self.model
144+
stage_manager = self.pipeline_stage_manager
145+
held_layers = []
146+
layers_per_stage = self.distribute_layers(len(module.h), stage_manager.num_stages)
147+
if self.stage_manager.is_first_stage():
148+
held_layers.append(module.word_embeddings)
149+
held_layers.append(module.word_embeddings_layernorm)
150+
151+
start_idx, end_idx = self.get_stage_index(layers_per_stage, self.stage_manager.stage)
152+
held_layers.extend(module.h[start_idx:end_idx])
153+
154+
if self.stage_manager.is_last_stage():
155+
held_layers.append(module.ln_f)
156+
157+
return held_layers
158+
159+
def get_shared_params(self, module: BloomModel) -> List[Dict[int, Tensor]]:
160+
'''no shared params in bloommodel'''
161+
pass
114162

115163

116164
class BloomForCausalLMPolicy(BloomPolicy):
@@ -181,3 +229,173 @@ def module_policy(self):
181229
class BloomForQuestionAnsweringPolicy(BloomPolicy):
182230
# No head sharding as the output features is only 2
183231
pass
232+
233+
234+
def bloom_model_forward(
235+
self: BloomModel,
236+
input_ids: Optional[torch.LongTensor] = None,
237+
past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
238+
attention_mask: Optional[torch.Tensor] = None,
239+
head_mask: Optional[torch.LongTensor] = None,
240+
inputs_embeds: Optional[torch.LongTensor] = None,
241+
use_cache: Optional[bool] = None,
242+
output_attentions: Optional[bool] = None,
243+
output_hidden_states: Optional[bool] = None,
244+
return_dict: Optional[bool] = None,
245+
stage_manager: Optional[PipelineStageManager] = None,
246+
hidden_states: Optional[torch.FloatTensor] = None,
247+
**deprecated_arguments,
248+
) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]:
249+
if deprecated_arguments.pop("position_ids", False) is not False:
250+
# `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None`
251+
warnings.warn(
252+
"`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore"
253+
" passing `position_ids`.",
254+
FutureWarning,
255+
)
256+
if len(deprecated_arguments) > 0:
257+
raise ValueError(f"Got unexpected arguments: {deprecated_arguments}")
258+
259+
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
260+
output_hidden_states = (output_hidden_states
261+
if output_hidden_states is not None else self.config.output_hidden_states)
262+
use_cache = use_cache if use_cache is not None else self.config.use_cache
263+
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
264+
265+
# add warnings here
266+
if output_attentions:
267+
logger.warning_once('output_attentions=True is not supported for pipeline models at the moment.')
268+
output_attentions = False
269+
if output_hidden_states:
270+
logger.warning_once('output_hidden_states=True is not supported for pipeline models at the moment.')
271+
output_hidden_states = False
272+
if use_cache:
273+
logger.warning_once('use_cache=True is not supported for pipeline models at the moment.')
274+
use_cache = False
275+
# Prepare head mask if needed
276+
# 1.0 in head_mask indicate we keep the head
277+
# attention_probs has shape batch_size x num_heads x N x N
278+
279+
# head_mask has shape n_layer x batch x num_heads x N x N
280+
head_mask = self.get_head_mask(head_mask, self.config.n_layer)
281+
282+
# case: First stage of training
283+
if stage_manager.is_first_stage():
284+
# check input_ids and inputs_embeds
285+
if input_ids is not None and inputs_embeds is not None:
286+
raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
287+
elif input_ids is not None:
288+
batch_size, seq_length = input_ids.shape
289+
elif inputs_embeds is not None:
290+
batch_size, seq_length, _ = inputs_embeds.shape
291+
else:
292+
raise ValueError("You have to specify either input_ids or inputs_embeds")
293+
294+
if inputs_embeds is None:
295+
inputs_embeds = self.word_embeddings(input_ids)
296+
297+
hidden_states = self.word_embeddings_layernorm(inputs_embeds)
298+
# initialize in the first stage and then pass to the next stage
299+
else:
300+
input_shape = hidden_states.shape[:-1]
301+
batch_size, seq_length = input_shape
302+
303+
# extra recording tensor should be generated in the first stage
304+
305+
presents = () if use_cache else None
306+
all_self_attentions = () if output_attentions else None
307+
all_hidden_states = () if output_hidden_states else None
308+
309+
if self.gradient_checkpointing and self.training:
310+
if use_cache:
311+
logger.warning_once(
312+
"`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`...")
313+
use_cache = False
314+
315+
if past_key_values is None:
316+
past_key_values = tuple([None] * len(self.h))
317+
# Compute alibi tensor: check build_alibi_tensor documentation,build for every stage
318+
seq_length_with_past = seq_length
319+
past_key_values_length = 0
320+
if past_key_values[0] is not None:
321+
past_key_values_length = past_key_values[0][0].shape[2] # source_len
322+
323+
seq_length_with_past = seq_length_with_past + past_key_values_length
324+
if attention_mask is None:
325+
attention_mask = torch.ones((batch_size, seq_length_with_past), device=hidden_states.device)
326+
else:
327+
attention_mask = attention_mask.to(hidden_states.device)
328+
329+
alibi = self.build_alibi_tensor(attention_mask, self.num_heads, dtype=hidden_states.dtype)
330+
331+
# causal_mask is constructed every stage and its input is passed through different stages
332+
causal_mask = self._prepare_attn_mask(
333+
attention_mask,
334+
input_shape=(batch_size, seq_length),
335+
past_key_values_length=past_key_values_length,
336+
)
337+
338+
# calculate the num_layers
339+
num_layers_per_stage = len(self.h) // stage_manager.num_stages
340+
start_layer = stage_manager.stage * num_layers_per_stage
341+
end_layer = (stage_manager.stage + 1) * num_layers_per_stage
342+
343+
for i, (block, layer_past) in enumerate(zip(self.h[start_layer:end_layer], past_key_values[start_layer:end_layer])):
344+
if output_hidden_states:
345+
all_hidden_states = all_hidden_states + (hidden_states,)
346+
347+
if self.gradient_checkpointing and self.training:
348+
349+
def create_custom_forward(module):
350+
351+
def custom_forward(*inputs):
352+
# None for past_key_value
353+
return module(*inputs, use_cache=use_cache, output_attentions=output_attentions)
354+
355+
return custom_forward
356+
357+
outputs = torch.utils.checkpoint.checkpoint(
358+
create_custom_forward(block),
359+
hidden_states,
360+
alibi,
361+
causal_mask,
362+
layer_past,
363+
head_mask[i],
364+
)
365+
else:
366+
outputs = block(
367+
hidden_states,
368+
layer_past=layer_past,
369+
attention_mask=causal_mask,
370+
head_mask=head_mask[i],
371+
use_cache=use_cache,
372+
output_attentions=output_attentions,
373+
alibi=alibi,
374+
)
375+
376+
hidden_states = outputs[0]
377+
378+
if use_cache is True:
379+
presents = presents + (outputs[1],)
380+
if output_attentions:
381+
all_self_attentions = all_self_attentions + \
382+
(outputs[2 if use_cache else 1],)
383+
384+
if stage_manager.is_last_stage():
385+
# Add last hidden state
386+
hidden_states = self.ln_f(hidden_states)
387+
388+
# TODO: deal with all_hidden_states, all_self_attentions, presents
389+
if output_hidden_states:
390+
all_hidden_states = all_hidden_states + (hidden_states,)
391+
392+
if not return_dict:
393+
return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)
394+
395+
# attention_mask is not returned ; presents = past_key_values
396+
return BaseModelOutputWithPastAndCrossAttentions(
397+
last_hidden_state=hidden_states,
398+
past_key_values=presents,
399+
hidden_states=all_hidden_states,
400+
attentions=all_self_attentions,
401+
)

0 commit comments

Comments
 (0)