pc and remove prints

mosaicml · Apr 12, 2024 · 5938048 · 5938048
1 parent 7c3d493
commit 5938048
Show file tree

Hide file tree

Showing 3 changed files with 1 addition and 4 deletions.
diff --git a/llmfoundry/models/layers/blocks.py b/llmfoundry/models/layers/blocks.py
@@ -61,7 +61,6 @@ def __init__(
         use_pad_tok_in_ffn: bool = True,
         **kwargs: Any,
     ):
-        print(ffn_config)
         if attn_config is None:
             attn_config = attn_config_defaults
 

diff --git a/llmfoundry/models/layers/layer_builders.py b/llmfoundry/models/layers/layer_builders.py
@@ -48,7 +48,7 @@ def build_ffn(
         'expansion_ratio': expansion_ratio,
         'device': device,
         'bias': bias,
-        **{k:v for k,v in ffn_kwargs.items() if k != 'ffn_type'},
+        **{k: v for k, v in ffn_kwargs.items() if k != 'ffn_type'},
     }
 
     def _validation_function(maybe_module: Any):

diff --git a/llmfoundry/models/mpt/modeling_mpt.py b/llmfoundry/models/mpt/modeling_mpt.py
@@ -325,7 +325,6 @@ def __init__(self, config: MPTConfig):
         self.emb_drop = nn.Dropout(config.emb_pdrop)
         self.mb_args = None
         block_args = config.to_dict()
-        print(block_args['ffn_config'], block_args['ffn_config']['ffn_type'])
         if block_args['ffn_config']['ffn_type'] in ffns_with_megablocks:
             block_args['ffn_config'] = config_moe_args(
                 block_args['ffn_config'],
@@ -334,7 +333,6 @@ def __init__(self, config: MPTConfig):
                 config.n_layers,
             )
             self.mb_args = block_args['ffn_config'].get('args')
-            print(block_args['ffn_config'], block_args['ffn_config']['ffn_type'])
 
         self.blocks = nn.ModuleList([
             MPTBlock(