[ROCm]: fix: reduce MoE temp memory — embedding cap, weight sum default, skip trivial specs (PR3)

[cj401-amd]

Summary

• Embeddings: cap use_iota_embed to ≤2 GiB one-hot size to prevent OOM on large
  vocabularies; add explicit nn.with_logical_constraint after embedding lookup
• MoE config: change float32_weight_sum default from true to false — the f32
  upcast adds ~2 GB temp per device with minimal numerical benefit for most configs
• DeepSeek: fix activation PartitionSpec to include fsdp_transpose and context
  axes; use remove_size_one_mesh_axis helper; remove redundant jax.reshard calls
• Mixtral: replace nn.with_logical_constraint with maybe_shard_with_logical(..., skip_trivial_specs=True) throughout MixtralDecoderLayer to avoid no-op sharding
  constraints that add XLA overhead

Test plan

• python3 -m pytest tests/unit/train_compile_test.py -v -k "moe or deepseek or mixtral"
• Smoke-test MoE model (e.g. mixtral-8x7b or deepseek3-test config)

[codecov]

Codecov Report

❌ Patch coverage is 53.78151% with 55 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/layers/pipeline.py	38.59%	35 Missing ⚠️
src/maxtext/layers/attention_op.py	0.00%	7 Missing ⚠️
src/maxtext/layers/normalizations.py	54.54%	4 Missing and 1 partial ⚠️
src/maxtext/trainers/pre_train/train.py	42.85%	3 Missing and 1 partial ⚠️
src/maxtext/models/mixtral.py	76.92%	0 Missing and 3 partials ⚠️
src/maxtext/models/deepseek.py	50.00%	1 Missing ⚠️

📢 Thoughts on this report? Let us know!