MCPcopy Create free account
hub / github.com/THUDM/LongWriter / forward

Method forward

train/patch/modeling_llama.py:150–170  ·  view source on GitHub ↗
(self, x, position_ids)

Source from the content-addressed store, hash-verified

148
149 @torch.no_grad()
150 def forward(self, x, position_ids):
151 if "dynamic" in self.rope_type:
152 self._dynamic_frequency_update(position_ids, device=x.device)
153
154 # Core RoPE block
155 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
156 position_ids_expanded = position_ids[:, None, :].float()
157 # Force float32 (see https://github.com/huggingface/transformers/pull/29285)
158 device_type = x.device.type
159 device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
160 with torch.autocast(device_type=device_type, enabled=False):
161 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
162 emb = torch.cat((freqs, freqs), dim=-1)
163 cos = emb.cos()
164 sin = emb.sin()
165
166 # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention
167 cos = cos * self.attention_scaling
168 sin = sin * self.attention_scaling
169
170 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
171
172
173class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding):

Callers

nothing calls this directly

Calls 1

Tested by

no test coverage detected