enable flexible emb init

examples/llm/src/models/param_init_fns.py

@@ -3,6 +3,7 @@
 
 import math
 import warnings
+from collections.abc import Sequence
 from functools import partial
 
 import torch
@@ -68,7 +69,38 @@ def generic_param_init_fn_(module, cfg, init_fn_):
 
     elif isinstance(module, nn.Embedding):
         # Embedding
-        init_fn_(module.weight)
+        if cfg.get('emb_init_std') is not None:
+            std = cfg.get('emb_init_std')
+            if std == 0:
+                warnings.warn(f'Embedding layer initialized to 0.')
+            emb_init_fn_ = partial(torch.nn.init.normal_, mean=0.0, std=std)
+            if cfg.get('verbose', 0) > 1:
+                warnings.warn(
+                    f'Embedding layer initialized using normal distribution with mean=0 and {std=}.'
+                )
+        elif cfg.get('emb_init_uniform_lim') is not None:
+            lim = cfg.get('emb_init_uniform_lim')
+            if isinstance(lim, Sequence):
+                if len(lim) > 2:
+                    raise ValueError(
+                        f'Uniform init requires a min and a max limit. User input: {lim}.'
+                    )
+                if lim[0] == lim[1]:
+                    warnings.warn(f'Embedding layer initialized to {lim[0]}.')
+            else:
+                if lim == 0:
+                    warnings.warn(f'Embedding layer initialized to 0.')
+                lim = [-lim, lim]
+            a, b = lim
+            emb_init_fn_ = partial(torch.nn.init.uniform_, a=a, b=b)
+            if cfg.get('verbose', 0) > 1:
+                warnings.warn(
+                    f'Embedding layer initialized using uniform distribution in range {lim}.'
+                )
+        else:
+            emb_init_fn_ = init_fn_
+
+        emb_init_fn_(module.weight)
 
     elif isinstance(module, nn.LayerNorm):
         # LayerNorm