(self)
| 48 | return self.state_dict(destination, prefix, keep_vars) |
| 49 | |
| 50 | def word_embeddings_weight(self): |
| 51 | if mpu.is_pipeline_first_stage(ignore_virtual=True): |
| 52 | return self.language_model.embedding.word_embeddings.weight |
| 53 | if mpu.is_pipeline_last_stage(ignore_virtual=True): |
| 54 | if not self.share_word_embeddings: |
| 55 | raise Exception( |
| 56 | "word_embeddings_weight() called for last " |
| 57 | "stage, but share_word_embeddings is false" |
| 58 | ) |
| 59 | return self.word_embeddings.weight |
| 60 | raise Exception( |
| 61 | "word_embeddings_weight() should be " "called for first and last stage only" |
| 62 | ) |
| 63 | |
| 64 | def initialize_word_embeddings(self, init_method_normal): |
| 65 | args = get_args() |
no outgoing calls
no test coverage detected