Use an env variable to roll out model sharding cache

Use an env variable to roll out model sharding cache #6015