feat: Add TP to embed_tokens and lm_head for Gemma models#879

terrykong merged 1 commit intomainNVIDIA-NeMo/RL:mainfrom

ruit/add_tp_for_gemmaNVIDIA-NeMo/RL:ruit/add_tp_for_gemmaCopy head branch name to clipboard