workable code (#11739)

* workable code * not set relative_attention_num_buckets, just use default value * remove relative_attention_num_buckets to use default values * fix minor * fix minor * fix pylint * Apply isort and black reformatting Signed-off-by: huvunvidia <[email protected]> --------- Signed-off-by: huvunvidia <[email protected]> Co-authored-by: Huy Vu2 <[email protected]> Co-authored-by: huvunvidia <[email protected]>
NVIDIA · Jan 5, 2025 · 638d0f7 · 638d0f7
1 parent 98f0b76
commit 638d0f7
Show file tree

Hide file tree

Showing 4 changed files with 518 additions and 10 deletions.
diff --git a/nemo/collections/llm/__init__.py b/nemo/collections/llm/__init__.py
@@ -134,7 +134,15 @@
 from nemo.collections.llm.t5.data import MockDataModule as T5MockDataModule
 from nemo.collections.llm.t5.data import PreTrainingDataModule as T5PreTrainingDataModule
 from nemo.collections.llm.t5.data import SquadDataModule as T5SquadDataModule
-from nemo.collections.llm.t5.model import T5Config, T5Model, t5_data_step, t5_forward_step
+from nemo.collections.llm.t5.model import (
+    T5Config,
+    T5Config3B,
+    T5Config11B,
+    T5Config220M,
+    T5Model,
+    t5_data_step,
+    t5_forward_step,
+)
 
 __all__ = [
     "MockDataModule",
@@ -145,6 +153,9 @@
     "gpt_forward_step",
     "T5Model",
     "T5Config",
+    "T5Config220M",
+    "T5Config3B",
+    "T5Config11B",
     "BertConfig",
     "BertModel",
     "t5_data_step",

diff --git a/nemo/collections/llm/t5/model/__init__.py b/nemo/collections/llm/t5/model/__init__.py
@@ -1,6 +1,9 @@
 from nemo.collections.llm.t5.model.t5 import (
     MaskedTokenLossReduction,
     T5Config,
+    T5Config3B,
+    T5Config11B,
+    T5Config220M,
     T5Model,
     local_layer_spec,
     t5_data_step,
@@ -10,6 +13,9 @@
 
 __all__ = [
     "T5Config",
+    "T5Config220M",
+    "T5Config3B",
+    "T5Config11B",
     "T5Model",
     "MaskedTokenLossReduction",
     "t5_data_step",