Closed FanhaiLu1 closed 3 months ago
This PR do two things: 1: Gemma replica sharding be consistent as llama 2: Set bfloat16 or float32 in make_env_tiny
This PR do two things: 1: Gemma replica sharding be consistent as llama 2: Set bfloat16 or float32 in make_env_tiny