Skip to content

Add accelerate support for M2M100 #19912

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
younesbelkada merged 2 commits into from
Oct 27, 2022
Merged

Add accelerate support for M2M100 #19912

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f863d35
add `accelerate` support for M2M100
younesbelkada Oct 26, 2022
208e9e6
fix device set nit
younesbelkada Oct 27, 2022
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
15 changes: 9 additions & 6 deletions src/transformers/models/m2m_100/modeling_m2m_100.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -532,6 +532,7 @@ class M2M100PreTrainedModel(PreTrainedModel):
config_class = M2M100Config
base_model_prefix = "model"
supports_gradient_checkpointing = True
_no_split_modules = ["M2M100Attention"]

def _init_weights(self, module):
std = self.config.init_std
Expand Down Expand Up @@ -693,10 +694,10 @@ def __init__(self, config: M2M100Config, embed_tokens: Optional[nn.Embedding] =
self.max_source_positions = config.max_position_embeddings
self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0

self.embed_tokens = nn.Embedding(config.vocab_size, embed_dim, self.padding_idx)

if embed_tokens is not None:
self.embed_tokens = embed_tokens
else:
self.embed_tokens = nn.Embedding(config.vocab_size, embed_dim, self.padding_idx)
self.embed_tokens.weight = embed_tokens.weight

self.embed_positions = M2M100SinusoidalPositionalEmbedding(
config.max_position_embeddings,
Expand Down Expand Up @@ -777,6 +778,7 @@ def forward(
inputs_embeds = self.embed_tokens(input_ids) * self.embed_scale

embed_pos = self.embed_positions(input_ids, inputs_embeds)
embed_pos = embed_pos.to(inputs_embeds.device)

hidden_states = inputs_embeds + embed_pos
hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
Expand Down Expand Up @@ -868,10 +870,10 @@ def __init__(self, config: M2M100Config, embed_tokens: Optional[nn.Embedding] =
self.max_target_positions = config.max_position_embeddings
self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0

self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model, self.padding_idx)

if embed_tokens is not None:
self.embed_tokens = embed_tokens
else:
self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model, self.padding_idx)
self.embed_tokens.weight = embed_tokens.weight

self.embed_positions = M2M100SinusoidalPositionalEmbedding(
config.max_position_embeddings,
Expand Down Expand Up @@ -1010,6 +1012,7 @@ def forward(

# embed positions
positions = self.embed_positions(input_ids, inputs_embeds, past_key_values_length)
positions = positions.to(inputs_embeds.device)

hidden_states = inputs_embeds + positions

Expand Down