Changes

• llama started giving compilation failures on pytorch 2.10 when running the model on TP > 1. There seems to be changes in torch 2.10.0, which made the tracking of fx graph more strict. In current implementation, we were wrapping TP strategy around embedding and head, which seemed not necessary, since llama like granite, doesn't use tied weights and doesn't need special handling for head or embedding. This likely created a differences in the graph and resulted in symbolic shape issues. We are already applying distributed strategy around individual Llama blocks.
• This PR simplifies this by removing those conditional blocks and this makes the model run on torch 2.10.