Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Phi3DecoderLayer`] We only modified the attention mask Args: config: Phi3Config
| 18 | |
| 19 | |
| 20 | class Phi3Transformer(Phi3Model): |
| 21 | """ |
| 22 | Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Phi3DecoderLayer`] |
| 23 | We only modified the attention mask |
| 24 | Args: |
| 25 | config: Phi3Config |
| 26 | """ |
| 27 | def prefetch_layer(self, layer_idx: int, device: torch.device): |
| 28 | "Starts prefetching the next layer cache" |
| 29 | with torch.cuda.stream(self.prefetch_stream): |
| 30 | # Prefetch next layer tensors to GPU |
| 31 | for name, param in self.layers[layer_idx].named_parameters(): |
| 32 | param.data = param.data.to(device, non_blocking=True) |
| 33 | |
| 34 | def evict_previous_layer(self, layer_idx: int): |
| 35 | "Moves the previous layer cache to the CPU" |
| 36 | prev_layer_idx = layer_idx - 1 |
| 37 | for name, param in self.layers[prev_layer_idx].named_parameters(): |
| 38 | param.data = param.data.to("cpu", non_blocking=True) |
| 39 | |
| 40 | def get_offlaod_layer(self, layer_idx: int, device: torch.device): |
| 41 | # init stream |
| 42 | if not hasattr(self, "prefetch_stream"): |
| 43 | self.prefetch_stream = torch.cuda.Stream() |
| 44 | |
| 45 | # delete previous layer |
| 46 | torch.cuda.current_stream().synchronize() |
| 47 | self.evict_previous_layer(layer_idx) |
| 48 | |
| 49 | # make sure the current layer is ready |
| 50 | torch.cuda.synchronize(self.prefetch_stream) |
| 51 | |
| 52 | # load next layer |
| 53 | self.prefetch_layer((layer_idx + 1) % len(self.layers), device) |
| 54 | |
| 55 | |
| 56 | def forward( |
| 57 | self, |
| 58 | input_ids: torch.LongTensor = None, |
| 59 | attention_mask: Optional[torch.Tensor] = None, |
| 60 | position_ids: Optional[torch.LongTensor] = None, |
| 61 | past_key_values: Optional[List[torch.FloatTensor]] = None, |
| 62 | inputs_embeds: Optional[torch.FloatTensor] = None, |
| 63 | use_cache: Optional[bool] = None, |
| 64 | output_attentions: Optional[bool] = None, |
| 65 | output_hidden_states: Optional[bool] = None, |
| 66 | return_dict: Optional[bool] = None, |
| 67 | cache_position: Optional[torch.LongTensor] = None, |
| 68 | offload_model: Optional[bool] = False, |
| 69 | ) -> Union[Tuple, BaseModelOutputWithPast]: |
| 70 | output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions |
| 71 | output_hidden_states = ( |
| 72 | output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states |
| 73 | ) |
| 74 | use_cache = use_cache if use_cache is not None else self.config.use_cache |
| 75 | |
| 76 | return_dict = return_dict if return_dict is not None else self.config.use_return_dict |
| 77 |