Fix some things that broke in the transformers (LLNL#2106)

* Fix some things that broke from the recent API changes * cleanup some scaling_factor=str() issues * Turns out, extend doesn't return the list :/
graham63 · May 2, 2022 · e9a31b5 · e9a31b5
1 parent 01cc59f
commit e9a31b5
Show file tree

Hide file tree

Showing 8 changed files with 13 additions and 13 deletions.
diff --git a/applications/nlp/transformer/subgraph/train.py b/applications/nlp/transformer/subgraph/train.py
@@ -66,7 +66,7 @@ def make_model(
     )
     embeddings = lbann.WeightedSum(
         embeddings,
-        scaling_factors=str(math.sqrt(embed_dim)),
+        scaling_factors=math.sqrt(embed_dim),
     )
     embeddings_slice = lbann.Slice(
         embeddings,

diff --git a/applications/nlp/transformer/subgraph/train_infer.py b/applications/nlp/transformer/subgraph/train_infer.py
@@ -66,7 +66,7 @@ def make_model(
     )
     embeddings = lbann.WeightedSum(
         embeddings,
-        scaling_factors=str(math.sqrt(embed_dim)),
+        scaling_factors=math.sqrt(embed_dim),
     )
     embeddings_slice = lbann.Slice(
         embeddings,

diff --git a/applications/nlp/transformer/train.py b/applications/nlp/transformer/train.py
@@ -56,7 +56,7 @@ def make_model(
     )
     embeddings = lbann.WeightedSum(
         embeddings,
-        scaling_factors=str(math.sqrt(embed_dim)),
+        scaling_factors=math.sqrt(embed_dim),
     )
     embeddings_slice = lbann.Slice(
         embeddings,

diff --git a/applications/physics/cosmology/ExaGAN/ExaGAN.py b/applications/physics/cosmology/ExaGAN/ExaGAN.py
@@ -92,7 +92,7 @@ def inv_transform(self,y):
                                       lbann.SafeDivide(
                                       lbann.Add(lbann.Constant(value=1.0, hint_layer=y),lbann.Identity(y)),
                                       lbann.Subtract(lbann.Constant(value=1.0, hint_layer=y),lbann.Identity(y))),
-                                      scaling_factors=str(self.datascale))
+                                      scaling_factors=self.datascale)
         linear_scale = 1/self.linear_scaler
-        CH2 = lbann.Tanh(lbann.WeightedSum(inv_transform,scaling_factors=str(linear_scale)))
+        CH2 = lbann.Tanh(lbann.WeightedSum(inv_transform,scaling_factors=linear_scale))
         return CH2
diff --git a/python/lbann/models/subgraph/transformer.py b/python/lbann/models/subgraph/transformer.py
@@ -1416,7 +1416,7 @@ def _subsequent_mask(self, size):
 		if size not in self._subsequent_mask_cache:
 			vals = np.triu(np.full((size,size), -1e9), k=1)
 			weights = lbann.Weights(
-				initializer=lbann.ValueInitializer(values=np.nditer(vals)),
+				initializer=lbann.ValueInitializer(values=vals.flat),
 				optimizer=None,
 				name=f'{self.name}_mask{size}_weights',
 			)

diff --git a/python/lbann/models/transformer.py b/python/lbann/models/transformer.py
@@ -50,12 +50,12 @@ def forward(self, x):
         # Affine transform
         s = lbann.WeightsLayer(
             weights=self.weight,
-            dims=f'1 {self.normalized_shape}',
+            dims=[1] + list(make_iterable(self.normalized_shape)),
         )
         s = lbann.Tessellate(s, hint_layer=x)
         b = lbann.WeightsLayer(
             weights=self.bias,
-            dims=f'1 {self.normalized_shape}',
+            dims=[1] + list(make_iterable(self.normalized_shape)),
         )
         b = lbann.Tessellate(b, hint_layer=x)
         x = lbann.Add(lbann.Multiply(s,x), b)
@@ -441,7 +441,7 @@ def _subsequent_mask(self, size):
         if size not in self._subsequent_mask_cache:
             vals = np.triu(np.full((size,size), -1e9), k=1)
             weights = lbann.Weights(
-                initializer=lbann.ValueInitializer(values=np.nditer(vals)),
+                initializer=lbann.ValueInitializer(values=vals.flat),
                 optimizer=None,
                 name=f'{self.name}_mask{size}_weights',
             )

diff --git a/python/lbann/modules/subgraph/transformer.py b/python/lbann/modules/subgraph/transformer.py
@@ -190,7 +190,7 @@ def forward(self, queries, keys, values, mask=None):
             )
             y = lbann.WeightedSum(
                 y,
-                scaling_factors=str(1 / math.sqrt(self.head_dim)),
+                scaling_factors=1 / math.sqrt(self.head_dim),
                 name=f'{head_name}_scale',
             )
 
@@ -432,7 +432,7 @@ def forward(self, queries, keys, values, mask=None):
             )
             y = lbann.WeightedSum(
                 y,
-                scaling_factors=str(1 / math.sqrt(self.head_dim)),
+                scaling_factors=1 / math.sqrt(self.head_dim),
                 name=f'{head_name}_scale',
             )
 
@@ -766,7 +766,7 @@ def forward(self, queries, keys, values, mask=None):
             )
             y = lbann.WeightedSum(
                 y,
-                scaling_factors=str(1 / math.sqrt(self.head_dim)),
+                scaling_factors=1 / math.sqrt(self.head_dim),
                 name=f'{head_name}_scale',
             )
 

diff --git a/python/lbann/modules/transformer.py b/python/lbann/modules/transformer.py
@@ -151,7 +151,7 @@ def forward(self, queries, keys, values, mask=None):
             )
             y = lbann.WeightedSum(
                 y,
-                scaling_factors=str(1 / math.sqrt(self.head_dim)),
+                scaling_factors=1 / math.sqrt(self.head_dim),
                 name=f'{head_name}_scale',
             )
             if mask: