diff --git a/config/default_config.yml b/config/default_config.yml
index 53f1ffc582..f283da18a0 100644
--- a/config/default_config.yml
+++ b/config/default_config.yml
@@ -66,6 +66,10 @@ num_register_tokens: 0
 
 healpix_level: 5
 
+# Split the encoder's HEALPix/location dimension across this many consecutive
+# distributed ranks through local assimilation.
+encoder_spatial_parallel_size: 4
+
 # Use 2D RoPE instead of traditional global positional encoding
 # When True: uses 2D RoPE based on healpix cell coordinates (lat/lon)
 # When False: uses traditional pe_global positional encoding
diff --git a/config/encoder_spatial_parallel_4.yml b/config/encoder_spatial_parallel_4.yml
new file mode 100644
index 0000000000..a0b2aad090
--- /dev/null
+++ b/config/encoder_spatial_parallel_4.yml
@@ -0,0 +1,8 @@
+# (C) Copyright 2025 WeatherGenerator contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+
+# Four consecutive ranks share one batch and each process one quarter of the
+# HEALPix/location dimension through local assimilation.
+encoder_spatial_parallel_size: 4
diff --git a/config/encoder_spatial_parallel_8.yml b/config/encoder_spatial_parallel_8.yml
new file mode 100644
index 0000000000..63319485b9
--- /dev/null
+++ b/config/encoder_spatial_parallel_8.yml
@@ -0,0 +1,8 @@
+# (C) Copyright 2025 WeatherGenerator contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+
+# Eight consecutive ranks share one batch and each process one eighth of the
+# HEALPix/location dimension through local assimilation.
+encoder_spatial_parallel_size: 8
diff --git a/docs/encoder_spatial_parallelism.md b/docs/encoder_spatial_parallelism.md
new file mode 100644
index 0000000000..7dcde4a2fc
--- /dev/null
+++ b/docs/encoder_spatial_parallelism.md
@@ -0,0 +1,812 @@
+# Encoder spatial parallelism
+
+This document describes the encoder spatial-parallel implementation introduced by
+[pkuyyj/WeatherGenerator PR #1](https://github.com/pkuyyj/WeatherGenerator/pull/1).
+It covers the two feature commits:
+
+| Commit | Purpose |
+| --- | --- |
+| `15dcff5a` | Add HEALPix encoder spatial parallelism |
+| `db4b477d` | Build encoder inputs on rank-local HEALPix domains |
+
+Upstream evaluation changes and branch-synchronization-only changes in the PR history are
+outside the scope of this document.
+
+## Motivation
+
+WeatherGenerator receives several heterogeneous input streams. A stream can be globally
+regular, such as ERA5, or contain a variable number of observations, such as
+`METOP_ABC_AVHRR_IASI`. Before this change, every distributed rank embedded every source
+token and performed local assimilation for every HEALPix cell.
+
+Embedding and local assimilation are spatially independent until the local cell
+representations are projected into the global latent representation. The implementation
+therefore distributes these stages over a HEALPix domain:
+
+1. ranks in one spatial group consume the same training sample;
+2. every rank owns a disjoint set of complete HEALPix cells;
+3. each of the nine streams is filtered and tokenized independently for that local domain;
+4. embedding, local assimilation, and local-to-global projection operate on local data;
+5. fixed-size per-cell latent representations are gathered in global cell order;
+6. query aggregation and global assimilation continue with the reconstructed global tensor.
+
+This reduces source-token and local-encoder activation memory. It does not shard the complete
+model or the complete training step.
+
+## Terminology
+
+| Term | Meaning |
+| --- | --- |
+| `world_size` | Total number of distributed ranks |
+| `spatial_parallel_size` | Number of ranks cooperating on one encoder input |
+| spatial group | Consecutive ranks that consume the same sample and partition its HEALPix domain |
+| data-parallel rank | Index of a spatial group in the global job |
+| spatial rank | Rank index within a spatial group |
+| data HEALPix level | The configured `healpix_level`, normally level 5 |
+| packed tokens | Variable-length tensor containing only existing stream tokens |
+| dense cell tensor | Fixed-size tensor with one position for every owned HEALPix cell |
+
+For global rank `r` and spatial-parallel size `S`:
+
+```text
+data_parallel_rank = r // S
+spatial_rank       = r % S
+```
+
+Ranks `[gS, gS + 1, ..., gS + S - 1]` form spatial group `g`.
+
+For example, with eight total ranks and `S = 4`:
+
+```text
+spatial group 0: global ranks [0, 1, 2, 3]
+spatial group 1: global ranks [4, 5, 6, 7]
+```
+
+The four ranks in each group consume the same sample. The two groups remain data-parallel
+with respect to each other.
+
+## HEALPix partitioning
+
+### Number of cells
+
+At HEALPix level `L`, the number of cells is:
+
+```text
+N(L) = 12 × 4^L
+```
+
+At the default data level `L = 5`:
+
+```text
+N(5) = 12 × 4^5 = 12,288
+```
+
+### Contiguous equal ownership
+
+The configured data-level cells are divided directly into equal, consecutive rank-local
+intervals:
+
+```text
+cells_per_rank = N(L) / S
+
+cell_start = spatial_rank × cells_per_rank
+cell_end   = cell_start + cells_per_rank
+```
+
+`N(L)` must be divisible by the spatial-parallel size. For level 5 and four spatial ranks:
+
+| Spatial rank | Level-5 cells | Number of level-5 cells |
+| ---: | ---: | ---: |
+| 0 | `[0, 3072)` | 3072 |
+| 1 | `[3072, 6144)` | 3072 |
+| 2 | `[6144, 9216)` | 3072 |
+| 3 | `[9216, 12288)` | 3072 |
+
+Common group sizes produce:
+
+| Spatial ranks | Level-5 cells/rank |
+| ---: | ---: |
+| 4 | 3072 |
+| 8 | 1536 |
+| 16 | 768 |
+| 32 | 384 |
+
+### Why consecutive ranges are valid
+
+The input mapping uses nested HEALPix ordering:
+
+```python
+ang2pix(2**healpix_level, theta, phi, nest=True)
+```
+
+Nested ordering gives every data-level cell a stable integer ID. Assigning consecutive ID
+intervals means that concatenating rank-local tensors in spatial-rank order reconstructs the
+original global cell order. No coarser parent-cell constraint is required.
+
+## Distributed topology and process groups
+
+`get_encoder_spatial_parallel_size()` validates that:
+
+- the configured spatial size is at least one;
+- it does not exceed `world_size`;
+- `world_size` is divisible by the spatial size.
+
+`get_encoder_spatial_parallel_group()` creates process groups from consecutive global ranks.
+All global ranks create the groups in the same order, and each process caches the group that
+contains it.
+
+With spatial size one, no additional process group is created and the feature reduces to the
+single-rank compatibility behavior.
+
+## End-to-end data flow
+
+The source path can be summarized as:
+
+```text
+source readers
+    │
+    ├─ same sample on every rank in a spatial group
+    │
+    ▼
+map every source location to nested HEALPix cell ID
+    │
+    ├─ filter cell_start <= cell_id < cell_end independently for each stream
+    │
+    ▼
+construct local stream tokens and local per-cell token counts
+    │
+    ▼
+embed every stream independently on its owning rank
+    │
+    ▼
+scatter stream-major embeddings into local cell-major packed order
+    │
+    ▼
+local assimilation
+    │
+    ▼
+local-to-global projection
+    │
+    ▼
+restore fixed-size dense local-cell tensor
+    │
+    ▼
+autograd-aware all-gather in spatial-rank order
+    │
+    ▼
+global cell tensor [0, N(healpix_level))
+    │
+    ▼
+query aggregation and global assimilation
+```
+
+### What is local and what remains global
+
+| Stage/data | Spatially sharded? |
+| --- | --- |
+| Source storage reads | No; each spatial rank currently reads the same source sample |
+| Source coordinate-to-HEALPix mapping | Computed independently on every spatial rank |
+| Source filtering and token construction | Yes |
+| Per-stream embedding | Yes |
+| Local assimilation | Yes |
+| Local-to-global projection | Yes |
+| Dense per-cell latent after gather | No |
+| Query aggregation | No |
+| Global assimilation | No |
+| Target construction, prediction, and loss | No |
+| Model parameters/FSDP state | Controlled separately by FSDP |
+
+The implementation reduces GPU source-token and activation memory, but does not yet perform
+distributed source I/O. Raw source reads and the coordinate mapping are replicated inside a
+spatial group.
+
+## Rank-local stream construction
+
+### Filtering mask
+
+Every stream is treated independently. After converting its coordinates to HEALPix IDs, the
+rank-local point mask is:
+
+```python
+local_domain_mask = (cell_ids >= cell_start) & (cell_ids < cell_end)
+```
+
+`np.flatnonzero()` returns the original indices of retained points. The retained points are
+then grouped by local HEALPix cell. The output is a list of length
+`cell_end - cell_start`; empty cells remain present as empty index arrays.
+
+This happens for every input stream, including:
+
+- `METOP_ABC_AVHRR_IASI`;
+- `ERA5_in`;
+- `ERA5`;
+- `METEOSAT_SEVIRI_IR`;
+- `GOES_ABI_IR`;
+- `HIMAWARI_AHI_IR`;
+- `GOES_ABI_VIS`;
+- `HIMAWARI_AHI_VIS`;
+- `SurfaceCombined`.
+
+Each stream can have a different number of retained observations and tokens on a rank.
+
+### Local tokenization
+
+`TokenizerMasking` receives `source_cell_start` and `source_cell_end`. Source windows call
+`tokenize_space()` or `tokenize_spacetime()` with this interval. Target windows continue to
+use the global interval.
+
+For each local source cell:
+
+1. locations are ordered by colatitude (`theta`);
+2. ordered locations are split into patches of the stream's configured `token_size`;
+3. the final patch is padded when required;
+4. masking selects retained patches without changing their relative cell order;
+5. `source_tokens_lens` records the number of retained patches in every local cell.
+
+`StreamData` is initialized with:
+
+```text
+source HEALPix cells = local_num_healpix_cells
+target HEALPix cells = num_healpix_cells
+```
+
+Consequently, for level 5 and four-way spatial parallelism:
+
+```python
+batch.source_samples.tokens_lens.shape[-1] == 3072
+batch.target_samples.tokens_lens.shape[-1] == 12288
+```
+
+### Variable-length METOP observations
+
+METOP is a useful example because its number of observations varies by time and domain.
+
+Before tokenization, METOP-A, METOP-B, and METOP-C reader rows are concatenated in configured
+reader order. Within each reader, source rows retain their storage order unless
+`shuffle_source` is enabled.
+
+For multiple source windows, `StreamData` stores step 0 as the newest window, followed by older
+windows. This step order is retained when the embedding engine concatenates the stream inputs.
+
+Tokenization deliberately converts this raw row order into cell-major order:
+
+```text
+local HEALPix cell 0
+    locations ordered north-to-south
+    patch 0: up to 512 locations
+    patch 1: up to 512 locations
+    ...
+local HEALPix cell 1
+    ...
+```
+
+A tensor such as:
+
+```text
+source_tokens_cells[step].shape = (N_rank, 512, 30)
+```
+
+has:
+
+- a variable `N_rank`, the number of retained METOP patch tokens on this rank;
+- 512 padded locations per patch;
+- 30 encoded features per location.
+
+`N_rank` does not need to be equal across spatial ranks. The variable-length METOP tensor is
+never directly passed to a fixed-shape all-gather.
+
+The temporary raw-row indices used by tokenization are not retained after source construction.
+The implementation preserves equivalence with the original cell-major model order, but does
+not provide an inverse map from assimilated latents back to original METOP-A/B/C storage rows.
+
+## Embedding order
+
+The embedding engine processes streams in configuration order. For each stream, it concatenates
+source tokens over input steps and batch samples and applies the stream-specific embedding
+network.
+
+At this point, embeddings are stream-major. `get_scatter_idxs_vectorized()` uses
+`batch.tokens_lens` to scatter them into packed cell-major order. Within each
+`(input_step, sample, cell)` position, streams retain configuration order and each stream
+retains its per-cell patch order.
+
+The positional-encoding index is derived from the same per-cell counts. Therefore the packed
+embedding tensor, its cell lengths, and its positional encodings remain aligned.
+
+If a rank owns no observations from any stream for a sample, the embedding engine returns an
+empty tensor rather than failing. The rank must still participate in later synchronized model
+and distributed operations.
+
+## Local assimilation and empty domains
+
+`cell_lens_local` describes the packed local token tensor in:
+
+```text
+(input_step, sample, local_cell)
+```
+
+order. Local assimilation uses these lengths to derive cumulative packed-token boundaries.
+Cell boundaries are never inferred by evenly slicing the token dimension, because token counts
+vary by stream, time, cell, and rank.
+
+When spatial parallelism is enabled, one local rank shard is processed as one local-assimilation
+chunk. This keeps the number and order of calls to FSDP-wrapped local modules consistent across
+spatial ranks.
+
+A local domain can contain no observations. Skipping FSDP modules on that rank would cause
+ranks to enter FSDP collectives in different orders and can deadlock backward. For an empty
+chunk, the implementation therefore:
+
+1. creates a one-token zero-valued dummy input;
+2. calls the local assimilation engine;
+3. calls latent interpolation;
+4. calls the local-to-global adapter;
+5. multiplies the result by zero and attaches it to the real output.
+
+The zero dependency has no numerical effect, but preserves the forward/backward graph and
+ensures that FSDP hooks execute consistently on every spatial rank.
+
+## Local-to-global boundary and all-gather
+
+The gather occurs after:
+
+```text
+stream embedding
+    → local assimilation
+    → latent interpolation
+    → local-to-global projection
+```
+
+It occurs before:
+
+```text
+query aggregation
+    → global assimilation
+```
+
+This is the principal synchronization boundary of the feature.
+
+### Why variable-length source tensors can be gathered safely
+
+The variable-length source tensors are first projected into a fixed number of query latents per
+cell. Before gathering, every rank restores a dense local tensor:
+
+```text
+[
+    input_steps × samples,
+    local_num_healpix_cells,
+    local_queries_per_cell,
+    global_embedding_dimension,
+]
+```
+
+Non-empty cells receive the projected local result. Empty cells retain their initialized latent
+slot. All ranks in a spatial group therefore have the same gather shape even when their raw
+observation and patch-token counts differ substantially.
+
+The implementation uses `torch.distributed.nn.functional.all_gather`, not the non-autograd
+collective, so gradients propagate from global processing back into the owning rank's local
+encoder path.
+
+### Reconstruction order
+
+The spatial process group is created from increasing, consecutive global ranks. The gather
+returns tensors in process-group rank order:
+
+```text
+[spatial rank 0, spatial rank 1, ..., spatial rank S - 1]
+```
+
+Each rank's local tensor is already ordered by increasing cell ID within its consecutive range.
+Concatenating gathered tensors along the cell dimension therefore reconstructs:
+
+```text
+global cell 0, global cell 1, ..., global cell N - 1
+```
+
+For four ranks at level 5:
+
+```text
+gather result =
+    rank 0 cells [0, 3072)
+    + rank 1 cells [3072, 6144)
+    + rank 2 cells [6144, 9216)
+    + rank 3 cells [9216, 12288)
+```
+
+The rank-local `tokens_lens` tensors are gathered and concatenated along their cell dimension in
+the same rank order. The reconstructed global cell mask is therefore aligned with the dense
+global latent tensor.
+
+The implementation then packs non-empty global cells for query aggregation and restores the
+global dense layout afterward.
+
+## Ordering guarantees
+
+The feature guarantees model-equivalent cell ordering:
+
+1. nested HEALPix maps every source location to a deterministic global cell ID;
+2. each global cell belongs to exactly one spatial rank;
+3. rank-local cell lists are ordered by increasing global cell ID;
+4. local tokenization matches the corresponding slice of global tokenization;
+5. stream embeddings are deterministically scattered using `tokens_lens`;
+6. dense local latent position `j` maps to global cell `cell_start + j`;
+7. all-gather results are concatenated in spatial-rank order;
+8. gathered lengths and gathered latents use the same order.
+
+This does not mean that the original source storage-row order survives tokenization. The model's
+canonical order is cell-major, followed by per-cell location/patch order. That was already the
+order consumed by local assimilation before spatial parallelism.
+
+The local/global construction equivalence test builds the global cell lists and the four local
+cell lists independently, concatenates the local lists, and checks exact array equality for every
+cell.
+
+### Sort stability caveat
+
+On NumPy 2.x, cell grouping explicitly requests a stable sort. On NumPy 1.x, the compatibility
+fallback currently uses NumPy's default `argsort`. The subsequent PyTorch colatitude sort is
+stable, so ordinary observations remain deterministic, but exact relative ordering of duplicate
+cell IDs with identical colatitudes is not formally guaranteed on the NumPy 1.x fallback.
+
+If exact duplicate-location ordering is required across NumPy versions, use:
+
+```python
+np.argsort(local_cell_ids, kind="stable")
+```
+
+and retain an explicit raw observation identifier for inverse mapping.
+
+## Backward-compatible full-grid input
+
+The encoder accepts either:
+
+- a rank-local source batch with `local_num_healpix_cells`; or
+- a legacy global source batch with `num_healpix_cells`.
+
+For a local batch, the embedding result is already local.
+
+For a global batch, `select_packed_cell_shard()` selects complete cells from the packed token
+tensor using `cell_lens`. This fallback preserves compatibility, but the selection occurs after
+embedding, so it does not provide the embedding-memory reduction of rank-local source
+construction.
+
+`select_packed_cell_shard()`:
+
+1. reshapes flattened cell lengths into complete HEALPix grids;
+2. marks `[cell_start, cell_end)` in every input-step/sample row;
+3. expands the cell mask by each cell's variable token count;
+4. selects the corresponding packed tokens;
+5. returns local cell lengths in unchanged packed order.
+
+It validates the cell range and verifies that the lengths describe the actual packed tensor.
+
+## Data-parallel training semantics
+
+Spatial ranks consume the same sample and must not be counted as independent data-parallel
+replicas.
+
+The effective data-parallel world size is:
+
+```text
+data_parallel_world_size = world_size / encoder_spatial_parallel_size
+```
+
+The effective batch size is:
+
+```text
+effective_batch_size =
+    batch_size_per_spatial_group × data_parallel_world_size
+```
+
+The trainer uses this effective data-parallel size for:
+
+- total batch-size reporting;
+- learning-rate scaling;
+- scheduler construction;
+- mini-epoch and continuation accounting.
+
+The sampler similarly maps global ranks in the same spatial group to one data-parallel rank, so
+they receive the same workset and random seed. Loader-worker IDs and mini-epoch indices still
+differentiate independent workers and epochs.
+
+When present in a saved run, `encoder_spatial_parallel_size_original` records the historical
+spatial size so effective-batch calculations remain consistent. If the key is absent, the
+implementation defaults it to the current spatial size; continuation from a pre-feature run
+therefore requires checking this value explicitly.
+
+## Configuration
+
+The relevant option is:
+
+```yaml
+encoder_spatial_parallel_size: 4
+```
+
+`encoder_spatial_parallel_size` controls the number of ranks per spatial group.
+
+Ready-to-use overrides are provided:
+
+```text
+config/encoder_spatial_parallel_4.yml
+config/encoder_spatial_parallel_8.yml
+```
+
+For a four-rank job in which all ranks cooperate on one encoder sample:
+
+```yaml
+encoder_spatial_parallel_size: 4
+```
+
+The expected derived values are:
+
+```text
+world_size: 4
+data_parallel_world_size: 1
+local level-5 cells/rank: 3072
+```
+
+For an eight-rank job with two four-rank spatial groups:
+
+```text
+world_size: 8
+encoder_spatial_parallel_size: 4
+data_parallel_world_size: 2
+```
+
+### Continuing an existing run
+
+A continued run can inherit a saved configuration that predates this feature. Pass the spatial
+configuration explicitly when the intended current run should use spatial parallelism.
+
+Verify the effective, merged configuration rather than relying only on the repository default.
+In a four-rank spatial-only run, `data_parallel_world_size` must be 1, not 4.
+
+## Runtime verification
+
+### Startup log
+
+Every spatial rank logs its fine-cell range. For level 5 and four ranks, expect:
+
+```text
+Encoder spatial rank 0/4 constructs source HEALPix cells [0, 3072)
+Encoder spatial rank 1/4 constructs source HEALPix cells [3072, 6144)
+Encoder spatial rank 2/4 constructs source HEALPix cells [6144, 9216)
+Encoder spatial rank 3/4 constructs source HEALPix cells [9216, 12288)
+```
+
+### Tensor-shape checks
+
+For level 5 and four spatial ranks:
+
+```python
+assert batch.source_samples.tokens_lens.shape[-1] == 3072
+assert batch.target_samples.tokens_lens.shape[-1] == 12288
+```
+
+For a globally regular stream such as ERA5, each rank should retain approximately one quarter of
+the source tokens. Exact equality depends on grid conventions and masks.
+
+For regional or observational streams such as METOP and geostationary satellite streams, token
+counts can be strongly imbalanced. The important invariants are:
+
+- every retained source token belongs to the rank's cell interval;
+- no source token belongs to two ranks;
+- the union of rank-local cell intervals covers the full grid.
+
+### Memory checks
+
+Use PyTorch allocated-memory statistics when measuring the feature:
+
+```python
+torch.cuda.reset_peak_memory_stats()
+
+# Run one representative training iteration.
+
+peak_allocated = torch.cuda.max_memory_allocated() / 2**30
+peak_reserved = torch.cuda.max_memory_reserved() / 2**30
+```
+
+`nvidia-smi` primarily reflects CUDA memory reserved by PyTorch. Released tensor memory remains
+in the caching allocator and may not visibly decrease even when live tensor memory does.
+
+Compare:
+
+- peak allocated memory;
+- peak reserved memory;
+- per-stage peaks;
+- all spatial ranks, because token imbalance can make one domain more expensive than another.
+
+## Tests
+
+`tests/test_encoder_spatial_parallel.py` covers:
+
+| Test area | Invariant |
+| --- | --- |
+| Local construction equivalence | Concatenated local cell lists equal global construction cell-by-cell |
+| Invalid local range | Out-of-range cell intervals are rejected |
+| Packed-token selection | Complete cells are selected across multiple input-step/sample rows |
+| Coverage and gradients | Shards cover every packed token exactly once and preserve gradients |
+| Invalid packed ranges | Invalid cell intervals are rejected |
+| Distributed size validation | Spatial groups must divide the distributed world |
+
+Recommended validation commands are:
+
+```bash
+pytest -q tests/test_encoder_spatial_parallel.py
+./scripts/actions.sh lint
+./scripts/actions.sh unit-test
+```
+
+A multi-rank integration run remains necessary to validate FSDP collective ordering, runtime
+memory, and real-stream token balance.
+
+## Expected memory behavior
+
+Memory savings are not expected to be exactly `1 / spatial_parallel_size`.
+
+Memory that should decrease includes:
+
+- rank-local source token tensors transferred to the GPU;
+- stream-embedding activations;
+- local-assimilation activations;
+- local-to-global projection activations.
+
+Memory that remains replicated or becomes global includes:
+
+- model parameters, gradients, and optimizer state according to the FSDP configuration;
+- source-reader CPU data before local filtering;
+- target tensors and target-side computation;
+- the dense gathered global latent tensor;
+- query aggregation;
+- global assimilation;
+- forecast engine and prediction heads;
+- CUDA allocator cache.
+
+The rank with the most observations can determine the job's usable batch size. Equal numbers of
+HEALPix cells do not imply equal numbers of stream tokens.
+
+## Troubleshooting
+
+### Embedding memory does not decrease
+
+Check:
+
+```python
+batch.source_samples.tokens_lens.shape[-1]
+```
+
+If it is 12,288 at level 5 in a four-way run, the data pipeline constructed a global source
+batch. The encoder can still select a local shard after embedding, but embedding remains global.
+
+Verify:
+
+- `encoder_spatial_parallel_size: 4` is present in the effective run configuration;
+- `data_parallel_world_size: 1` for a four-rank spatial-only run;
+- all four ownership messages appear in the startup log;
+- each source stream has a rank-local token count.
+
+### `nvidia-smi` remains high
+
+Compare `torch.cuda.max_memory_allocated()` with `torch.cuda.max_memory_reserved()`. A large
+difference usually represents reusable cached allocator segments, not live model tensors.
+
+### A rank has no observations
+
+This is valid for regional streams and sparse samples. The embedding engine returns an empty
+packed tensor, and the local assimilation path executes a zero-valued dummy dependency to keep
+FSDP calls synchronized.
+
+### Gather shape mismatch
+
+The gather must receive:
+
+```text
+[steps × samples, local_cells, queries_per_cell, embedding_dimension]
+```
+
+with identical `local_cells`, query count, dtype, and embedding dimension on all ranks in the
+spatial group. Variable source-token counts must be resolved before this boundary.
+
+### Global ordering mismatch
+
+Verify together:
+
+- the identical contiguous local range computed by the sampler and encoder;
+- nested HEALPix mapping (`nest=True`);
+- local `tokens_lens` cell dimension;
+- rank order inside the spatial process group;
+- concatenation along the cell dimension;
+- identical ordering for gathered lengths and gathered latents.
+
+Do not compare assimilated output against raw source-reader row order. Compare it against the
+non-parallel model's cell-major token and latent order.
+
+### FSDP hang during backward
+
+Confirm that:
+
+- all ranks in a spatial group consume the same sample;
+- all ranks enter the local FSDP modules in the same order;
+- empty ranks execute the dummy dependency path;
+- `world_size` is divisible by the spatial size;
+- no rank independently skips an otherwise valid sample because its local source domain is empty.
+
+## File-level implementation map
+
+| File | Responsibility |
+| --- | --- |
+| `config/default_config.yml` | Default spatial size |
+| `config/encoder_spatial_parallel_4.yml` | Four-rank override |
+| `config/encoder_spatial_parallel_8.yml` | Eight-rank override |
+| `src/weathergen/datasets/healpix_domain.py` | Rank-local point grouping |
+| `src/weathergen/datasets/multi_stream_data_sampler.py` | Shared spatial-group samples, local ownership, and runtime logging |
+| `src/weathergen/datasets/stream_data.py` | Separate source-local and target-global cell counts |
+| `src/weathergen/datasets/tokenizer_masking.py` | Local source tokenization range |
+| `src/weathergen/datasets/tokenizer_utils.py` | Nested HEALPix mapping and local token construction |
+| `src/weathergen/model/engines.py` | Local per-stream embedding and empty-domain handling |
+| `src/weathergen/model/spatial_parallel.py` | Legacy packed-token shard selection |
+| `src/weathergen/model/encoder.py` | Local assimilation, local-to-global projection, and differentiable gather |
+| `src/weathergen/train/trainer.py` | Effective data-parallel batch and scheduler semantics |
+| `src/weathergen/utils/distributed.py` | Spatial group validation and construction |
+| `tests/test_encoder_spatial_parallel.py` | Cell ownership, ordering, coverage, validation, and gradient tests |
+
+## Commit-by-commit design history
+
+### `15dcff5a`: Add HEALPix encoder spatial parallelism
+
+This commit establishes the initial model-side and distributed design:
+
+- adds `encoder_spatial_parallel_size`;
+- creates four-rank and eight-rank configuration overrides;
+- creates consecutive-rank spatial process groups;
+- makes spatial-group ranks consume the same data sample;
+- changes effective batch and scheduler scaling from global world size to data-parallel group count;
+- introduces `select_packed_cell_shard()` for variable-length cell-packed tensors;
+- selects local cells before local assimilation;
+- builds local latent query seeds and local positional encodings;
+- performs local assimilation and local-to-global projection on each rank;
+- restores dense local cells and gathers them before global processing;
+- uses an autograd-aware gather;
+- adds synchronized handling for empty local domains;
+- adds initial unit tests for packed-cell coverage, gradients, and distributed-size validation.
+
+At this stage, embedding still runs on the global packed source tensor and the result is selected
+after embedding.
+
+### `db4b477d`: Build encoder inputs on local HEALPix domains
+
+This commit moves the domain boundary into the data pipeline:
+
+- introduces rank-local HEALPix point grouping;
+- passes source cell ranges into source tokenization;
+- constructs only rank-local source cells and source token tensors;
+- keeps targets global;
+- adds separate source and target HEALPix dimensions to `StreamData`;
+- makes embedding operate on local source tensors;
+- gathers local `tokens_lens` to reconstruct the global mask;
+- accepts both local and legacy global source batches;
+- handles ranks with no embedded streams;
+- validates local/global batch cell dimensions;
+- adds exact local-construction-versus-global-slice tests.
+
+This is the commit that enables embedding activation memory to scale with the local observation
+domain.
+
+## Review checklist
+
+Before merging or extending this feature, verify:
+
+- [ ] The effective spatial size is explicit in the run configuration.
+- [ ] `world_size % encoder_spatial_parallel_size == 0`.
+- [ ] The data-level HEALPix cell count is divisible by the spatial size.
+- [ ] Every spatial group consumes identical samples.
+- [ ] Every source stream constructs only local cells.
+- [ ] Source and target cell dimensions remain intentionally different.
+- [ ] Variable-length stream tokens are resolved into fixed-size cell latents before gather.
+- [ ] Local lengths and local latents are gathered in identical rank order.
+- [ ] The gather remains autograd-aware.
+- [ ] Empty local domains enter all FSDP modules in consistent order.
+- [ ] Effective batch size counts spatial groups, not individual spatial ranks.
+- [ ] Tests cover exact ordering, token coverage, and gradients.
+- [ ] Runtime validation compares allocated memory across all spatial ranks.
diff --git a/src/weathergen/datasets/healpix_domain.py b/src/weathergen/datasets/healpix_domain.py
new file mode 100644
index 0000000000..7d0a069e42
--- /dev/null
+++ b/src/weathergen/datasets/healpix_domain.py
@@ -0,0 +1,50 @@
+# (C) Copyright 2025 WeatherGenerator contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+import numpy as np
+
+
+def build_local_healpix_cell_splits(
+    cell_ids: np.typing.NDArray[np.integer],
+    num_cells: int,
+    cell_start: int,
+    cell_end: int,
+) -> list[np.typing.NDArray[np.int64]]:
+    """Group original point indices for one consecutive HEALPix-cell domain."""
+
+    if not 0 <= cell_start < cell_end <= num_cells:
+        raise ValueError(
+            f"invalid HEALPix cell range [{cell_start}, {cell_end}) for {num_cells} cells"
+        )
+
+    # Domain-parallel filtering mask: this is applied independently to every
+    # stream immediately after its coordinates have been mapped to nested
+    # HEALPix cell IDs.
+    local_domain_mask = (cell_ids >= cell_start) & (cell_ids < cell_end)
+    local_point_idxs = np.flatnonzero(local_domain_mask)
+    local_cell_ids = cell_ids[local_point_idxs]
+    cell_splits = [np.array([], dtype=np.int64) for _ in range(cell_end - cell_start)]
+    if local_point_idxs.size == 0:
+        return cell_splits
+
+    stable_args = {"stable": True} if int(np.__version__.split(".")[0]) >= 2 else {}
+    local_order = np.argsort(local_cell_ids, **stable_args)
+    sorted_point_idxs = local_point_idxs[local_order]
+    sorted_cell_ids = local_cell_ids[local_order]
+    split_offsets = np.flatnonzero(np.diff(sorted_cell_ids))
+    point_idxs_by_occupied_cell = np.split(sorted_point_idxs, split_offsets + 1)
+
+    for cell_id, point_idxs in zip(
+        np.unique(sorted_cell_ids),
+        point_idxs_by_occupied_cell,
+        strict=True,
+    ):
+        cell_splits[cell_id - cell_start] = point_idxs
+
+    return cell_splits
diff --git a/src/weathergen/datasets/multi_stream_data_sampler.py b/src/weathergen/datasets/multi_stream_data_sampler.py
index bd4a8ad87e..ed07cebe3e 100644
--- a/src/weathergen/datasets/multi_stream_data_sampler.py
+++ b/src/weathergen/datasets/multi_stream_data_sampler.py
@@ -34,7 +34,7 @@
 )
 from weathergen.readers_extra.registry import get_extra_reader
 from weathergen.train.utils import Stage, get_batch_size_from_config
-from weathergen.utils.distributed import is_root
+from weathergen.utils.distributed import get_encoder_spatial_parallel_size, is_root
 
 type AnyDataReader = DataReaderBase | DataReaderAnemoi | DataReaderObs
 type StreamName = str
@@ -100,15 +100,41 @@ def __init__(self, cf: Config, mode_cfg: dict, stage: Stage):
 
         self.mini_epoch = 0
         self.mask_value = 0.0
-        self.rank = cf.rank
-        self.world_size = cf.world_size
+        # Ranks in one encoder-spatial group must consume the same batch. Data
+        # parallelism therefore operates across groups, not across individual ranks.
+        spatial_parallel_size = get_encoder_spatial_parallel_size(cf)
+        self.spatial_parallel_size = spatial_parallel_size
+        self.spatial_parallel_rank = cf.rank % spatial_parallel_size
+        self.rank = cf.rank // spatial_parallel_size
+        self.world_size = cf.world_size // spatial_parallel_size
         self.repeat_data = cf.data_loading.get("repeat_data_in_mini_epoch", False)
 
         # initialise healpic
         self.healpix_level = cf.healpix_level
         self.num_healpix_cells = 12 * 4**self.healpix_level
+        if self.num_healpix_cells % spatial_parallel_size:
+            raise ValueError(
+                f"number of HEALPix cells ({self.num_healpix_cells}) must be divisible by "
+                f"encoder_spatial_parallel_size ({spatial_parallel_size})"
+            )
+        self.local_num_healpix_cells = self.num_healpix_cells // spatial_parallel_size
+        self.local_cell_start = self.spatial_parallel_rank * self.local_num_healpix_cells
+        self.local_cell_end = self.local_cell_start + self.local_num_healpix_cells
         self.masker = Masker(cf.healpix_level, stage, cf.streams, self.mode_cfg)
-        self.tokenizer = TokenizerMasking(cf.healpix_level, self.masker)
+        self.tokenizer = TokenizerMasking(
+            cf.healpix_level,
+            self.masker,
+            self.local_cell_start,
+            self.local_cell_end,
+        )
+        if spatial_parallel_size > 1:
+            logger.info(
+                "Encoder spatial rank %d/%d constructs source HEALPix cells [%d, %d)",
+                self.spatial_parallel_rank,
+                spatial_parallel_size,
+                self.local_cell_start,
+                self.local_cell_end,
+            )
 
         forecast_cfg = FORECAST_DEFAULTS | OmegaConf.to_object(mode_cfg.get("forecast", {}))
         self.output_offset = forecast_cfg["offset"]
@@ -532,6 +558,7 @@ def _build_stream_data(
             num_steps_input,
             num_output_steps,
             self.num_healpix_cells,
+            source_healpix_cells=self.local_num_healpix_cells,
         )
 
         stream_data = self._build_stream_data_input(
@@ -703,7 +730,9 @@ def _get_batch(self, idx: int, num_forecast_steps: int):
 
             # tokenize windows
             # *_tokens = [ (cells_idx, cells_idx_lens), ... ] with length = #time_steps
-            input_tokens = self.tokenizer.get_tokens_windows(stream_info, input_data, True)
+            input_tokens = self.tokenizer.get_tokens_windows(
+                stream_info, input_data, True, local_source=True
+            )
             output_tokens = self.tokenizer.get_tokens_windows(stream_info, output_data, False)
 
             for sidx, source_mask in enumerate(source_masks.masks):
@@ -791,7 +820,13 @@ def __iter__(self) -> ModelBatch:
                 # ensure the batch is valid, i.e. not completely empty and no NaN values
                 # student teacher has no classical targets
                 mode = self.mode_cfg.get("training_mode")
-                not_valid = batch.sources_empty() or batch.is_nan()
+                # A valid global sample may have no observations in one rank's
+                # local source domain. All spatial ranks must still emit the
+                # same sample and enter encoder collectives in lockstep.
+                local_sources_empty = batch.sources_empty()
+                not_valid = (
+                    local_sources_empty if self.spatial_parallel_size == 1 else False
+                ) or batch.is_nan()
                 not_valid = not_valid or (batch.targets_empty() if "masking" in mode else False)
 
                 # skip completely empty batch item or when all targets are empty -> no grad
@@ -822,12 +857,8 @@ def worker_workset(self):
             # happens for each mini_epoch, for train and validation, and independently for each DDP
             # worker. After the bit-wise copy, the rng seed needs to be made unique for
             # DDP workers, loader process, mini_epoch.
-            dist = torch.distributed
             self.data_loader_rng_seed *= (
-                (((dist.get_rank() + 1) * 73) if dist.is_initialized() else 1)
-                * ((worker_info.id + 1) * 37)
-                * (self.mini_epoch + 13)
-                * 7
+                ((self.rank + 1) * 73) * ((worker_info.id + 1) * 37) * (self.mini_epoch + 13) * 7
             )
             # split workload
             per_worker = (local_end - local_start) // worker_info.num_workers
diff --git a/src/weathergen/datasets/stream_data.py b/src/weathergen/datasets/stream_data.py
index 4c3bb02b8b..c5e6528b82 100644
--- a/src/weathergen/datasets/stream_data.py
+++ b/src/weathergen/datasets/stream_data.py
@@ -58,6 +58,7 @@ def __init__(
         input_steps: int,
         output_steps: int,
         healpix_cells: int,
+        source_healpix_cells: int | None = None,
     ) -> None:
         """
         StreamData object
@@ -72,7 +73,10 @@ def __init__(
             Number of output steps
             Note -- Last input step and first output step always overlap.
         healpix_cells : int
-            Number of healpix cells for source
+            Number of global healpix cells used for targets
+        source_healpix_cells : int | None
+            Number of rank-local healpix cells used for encoder sources. Defaults
+            to ``healpix_cells`` for non-spatial-parallel callers.
 
         Returns
         -------
@@ -84,6 +88,9 @@ def __init__(
         self.input_steps = input_steps
         self.output_steps = output_steps
         self.healpix_cells = healpix_cells
+        self.source_healpix_cells = (
+            healpix_cells if source_healpix_cells is None else source_healpix_cells
+        )
 
         self.source_is_spoof = [False for _ in range(self.input_steps)]
         self.target_is_spoof = [False for _ in range(self.output_steps)]
@@ -104,7 +111,8 @@ def __init__(
         self.source_tokens_cells = [None for _ in range(self.input_steps)]
         # length of source tokens per cell (without padding)
         self.source_tokens_lens = [
-            torch.zeros(self.healpix_cells, dtype=torch.int32) for _ in range(self.input_steps)
+            torch.zeros(self.source_healpix_cells, dtype=torch.int32)
+            for _ in range(self.input_steps)
         ]
         # unprocessed source (for logging)
         self.source_raw = [None for _ in range(self.input_steps)]
diff --git a/src/weathergen/datasets/tokenizer_masking.py b/src/weathergen/datasets/tokenizer_masking.py
index 7c033e398f..0491ef30d7 100644
--- a/src/weathergen/datasets/tokenizer_masking.py
+++ b/src/weathergen/datasets/tokenizer_masking.py
@@ -40,11 +40,28 @@ def readerdata_to_torch(rdata: IOReaderData) -> IOReaderData:
 
 
 class TokenizerMasking(Tokenizer):
-    def __init__(self, healpix_level: int, masker: Masker):
+    def __init__(
+        self,
+        healpix_level: int,
+        masker: Masker,
+        source_cell_start: int = 0,
+        source_cell_end: int | None = None,
+    ):
         super().__init__(healpix_level)
         self.masker = masker
         self.rng = None
         self.token_size = None
+        self.source_cell_start = source_cell_start
+        self.source_cell_end = (
+            self.num_healpix_cells_source if source_cell_end is None else source_cell_end
+        )
+        if not (
+            0 <= self.source_cell_start < self.source_cell_end <= self.num_healpix_cells_source
+        ):
+            raise ValueError(
+                f"invalid source HEALPix cell range "
+                f"[{self.source_cell_start}, {self.source_cell_end})"
+            )
 
     def reset_rng(self, rng) -> None:
         """
@@ -53,7 +70,7 @@ def reset_rng(self, rng) -> None:
         self.masker.reset_rng(rng)
         self.rng = rng
 
-    def get_tokens_windows(self, stream_info, data, pad_tokens):
+    def get_tokens_windows(self, stream_info, data, pad_tokens, local_source=False):
         """
         Tokenize data (to amortize over the different views that are generated)
 
@@ -63,6 +80,8 @@ def get_tokens_windows(self, stream_info, data, pad_tokens):
         tok = tokenize_spacetime if tok_spacetime else tokenize_space
         hl = self.healpix_level
         token_size = stream_info["token_size"]
+        cell_start = self.source_cell_start if local_source else 0
+        cell_end = self.source_cell_end if local_source else self.num_healpix_cells_source
 
         tokens = []
         for rdata in data:
@@ -72,7 +91,12 @@ def get_tokens_windows(self, stream_info, data, pad_tokens):
                 continue
             # tokenize data
             idxs_cells, idxs_cells_lens = tok(
-                readerdata_to_torch(rdata), token_size, hl, pad_tokens
+                readerdata_to_torch(rdata),
+                token_size,
+                hl,
+                pad_tokens,
+                cell_start=cell_start,
+                cell_end=cell_end,
             )
             tokens += [(idxs_cells, idxs_cells_lens)]
 
@@ -129,6 +153,7 @@ def get_source(
     ):
         # create tokenization index
         (idxs_cells, idxs_cells_lens) = idxs_cells_data
+        cell_mask = cell_mask[self.source_cell_start : self.source_cell_end]
 
         # select strategy from XXX depending on stream and if student or teacher
 
@@ -144,7 +169,7 @@ def get_source(
             stream_info["stream_id"],
             rdata,
             time_win,
-            self.hpy_verts_rots_source[-1],
+            self.hpy_verts_rots_source[-1][self.source_cell_start : self.source_cell_end],
             encode_times_source,
         )
 
diff --git a/src/weathergen/datasets/tokenizer_utils.py b/src/weathergen/datasets/tokenizer_utils.py
index 1bd1d1722d..571d2d9cbd 100644
--- a/src/weathergen/datasets/tokenizer_utils.py
+++ b/src/weathergen/datasets/tokenizer_utils.py
@@ -5,6 +5,7 @@
 from torch import Tensor
 
 from weathergen.common.io import IOReaderData
+from weathergen.datasets.healpix_domain import build_local_healpix_cell_splits
 from weathergen.datasets.utils import (
     locs_to_cell_coords_ctrs,
     locs_to_ctr_coords,
@@ -12,10 +13,6 @@
     s2tor3,
 )
 
-# on some clusters our numpy version is pinned to be 1.x.x where the np.argsort does not
-# the stable=True argument
-numpy_argsort_args = {"stable": True} if int(np.__version__.split(".")[0]) >= 2 else {}
-
 
 def theta_phi_to_standard_coords(coords):
     thetas = ((90.0 - coords[:, 0]) / 180.0) * np.pi
@@ -95,34 +92,46 @@ def encode_times_target(times, time_win) -> torch.tensor:
     return time_tensor + 0.5
 
 
-def hpy_cell_splits(coords: torch.tensor, hl: int):
+def hpy_cell_splits(
+    coords: torch.tensor,
+    hl: int,
+    cell_start: int = 0,
+    cell_end: int | None = None,
+):
     """Compute healpix cell id for each coordinate on given level hl
 
     Returns
-      hpy_idxs_ord_split : list of per cell indices into thetas,phis,posr3
+      hpy_idxs_ord_split : list of per-local-cell indices into thetas,phis,posr3
       thetas : thetas in rad
       phis : phis in rad
     """
+    num_healpix_cells = 12 * 4**hl
+    cell_end = num_healpix_cells if cell_end is None else cell_end
+
     thetas, phis = theta_phi_to_standard_coords(coords)
     # healpix cells for all points
     hpy_idxs = ang2pix(2**hl, thetas, phis, nest=True)
 
-    # extract information to split according to cells by first sorting and then finding split idxs
-    hpy_idxs_ord = np.argsort(hpy_idxs, **numpy_argsort_args)
-    splits = np.flatnonzero(np.diff(hpy_idxs[hpy_idxs_ord]))
-
-    # extract per cell data
-    hpy_idxs_ord_temp = np.split(hpy_idxs_ord, splits + 1)
-    hpy_idxs_ord_split = [np.array([], dtype=np.int64) for _ in range(12 * 4**hl)]
-    # TODO: split smarter (with a augmented splits list?) so that this loop is not needed
-    for b, x in zip(np.unique(np.unique(hpy_idxs[hpy_idxs_ord])), hpy_idxs_ord_temp, strict=True):
-        hpy_idxs_ord_split[b] = x
+    # Nested HEALPix IDs make a consecutive interval a complete rank-local
+    # domain. The helper applies the point mask and builds only local cells.
+    hpy_idxs_ord_split = build_local_healpix_cell_splits(
+        hpy_idxs,
+        num_healpix_cells,
+        cell_start,
+        cell_end,
+    )
 
     return (hpy_idxs_ord_split, thetas, phis)
 
 
 def hpy_splits(
-    coords: torch.Tensor, hl: int, token_size: int, pad_tokens: bool, offset_step: int = 0
+    coords: torch.Tensor,
+    hl: int,
+    token_size: int,
+    pad_tokens: bool,
+    offset_step: int = 0,
+    cell_start: int = 0,
+    cell_end: int | None = None,
 ) -> tuple[list[torch.Tensor], list[torch.Tensor], torch.Tensor]:
     """Compute healpix cell for each data point and splitting information per cell;
        when the token_size is exceeded then splitting based on lat is used;
@@ -135,7 +144,9 @@ def hpy_splits(
     """
 
     # list of data points per healpix cell
-    (hpy_idxs_ord_split, thetas, phis) = hpy_cell_splits(coords, hl)
+    (hpy_idxs_ord_split, thetas, phis) = hpy_cell_splits(
+        coords, hl, cell_start=cell_start, cell_end=cell_end
+    )
 
     # if token_size is exceeed split based on latitude
     # TODO: split by hierarchically traversing healpix scheme
@@ -179,11 +190,21 @@ def tokenize_space(
     hl,
     pad_tokens=True,
     offset_step=0,
+    cell_start=0,
+    cell_end=None,
 ):
     """Process one window into tokens"""
 
     # idx_ord_lens is length is number of tokens per healpix cell
-    idxs_ord, idxs_ord_lens = hpy_splits(rdata.coords, hl, token_size, pad_tokens, offset_step)
+    idxs_ord, idxs_ord_lens = hpy_splits(
+        rdata.coords,
+        hl,
+        token_size,
+        pad_tokens,
+        offset_step,
+        cell_start,
+        cell_end,
+    )
 
     return idxs_ord, idxs_ord_lens
 
@@ -193,14 +214,17 @@ def tokenize_spacetime(
     token_size,
     hl,
     pad_tokens=True,
+    cell_start=0,
+    cell_end=None,
 ):
     """Tokenize respecting an intrinsic time step in the data, i.e. each time step is tokenized
     separately
     """
 
     num_healpix_cells = 12 * 4**hl
-    idxs_cells = [[] for _ in range(num_healpix_cells)]
-    idxs_cells_lens = [[] for _ in range(num_healpix_cells)]
+    cell_end = num_healpix_cells if cell_end is None else cell_end
+    idxs_cells = [[] for _ in range(cell_end - cell_start)]
+    idxs_cells_lens = [[] for _ in range(cell_end - cell_start)]
 
     offset_step = 0
     t_unique = np.unique(rdata.datetimes)
@@ -210,7 +234,15 @@ def tokenize_spacetime(
         rdata_cur = IOReaderData(
             rdata.coords[mask], rdata.geoinfos[mask], rdata.data[mask], rdata.datetimes[mask]
         )
-        idxs_cur, idxs_cur_lens = tokenize_space(rdata_cur, token_size, hl, pad_tokens, offset_step)
+        idxs_cur, idxs_cur_lens = tokenize_space(
+            rdata_cur,
+            token_size,
+            hl,
+            pad_tokens,
+            offset_step,
+            cell_start,
+            cell_end,
+        )
 
         # collect data for all time steps
         idxs_cells = [t + tc for t, tc in zip(idxs_cells, idxs_cur, strict=True)]
diff --git a/src/weathergen/model/encoder.py b/src/weathergen/model/encoder.py
index 3e9f99a706..67cb92c4e5 100644
--- a/src/weathergen/model/encoder.py
+++ b/src/weathergen/model/encoder.py
@@ -9,6 +9,7 @@
 
 import torch
 from astropy_healpix import healpy
+from torch.distributed.nn.functional import all_gather
 from torch.utils.checkpoint import checkpoint
 
 from weathergen.common.config import Config
@@ -25,6 +26,11 @@
 # from weathergen.model.model import ModelParams
 from weathergen.model.parametrised_prob_dist import LatentInterpolator
 from weathergen.model.positional_encoding import positional_encoding_harmonic
+from weathergen.model.spatial_parallel import select_packed_cell_shard
+from weathergen.utils.distributed import (
+    get_encoder_spatial_parallel_group,
+    get_encoder_spatial_parallel_size,
+)
 
 
 class EncoderModule(torch.nn.Module):
@@ -43,6 +49,18 @@ def __init__(self, cf: Config, sources_size, targets_num_channels, targets_coord
 
         self.healpix_level = cf.healpix_level
         self.num_healpix_cells = 12 * 4**self.healpix_level
+        self.spatial_parallel_size = get_encoder_spatial_parallel_size(cf)
+        if self.num_healpix_cells % self.spatial_parallel_size:
+            raise ValueError(
+                f"number of HEALPix cells ({self.num_healpix_cells}) must be divisible by "
+                f"encoder_spatial_parallel_size ({self.spatial_parallel_size})"
+            )
+        self.spatial_parallel_group, self.spatial_parallel_rank = (
+            get_encoder_spatial_parallel_group(cf)
+        )
+        self.local_num_healpix_cells = self.num_healpix_cells // self.spatial_parallel_size
+        self.local_cell_start = self.spatial_parallel_rank * self.local_num_healpix_cells
+        self.local_cell_end = self.local_cell_start + self.local_num_healpix_cells
 
         self.cf = cf
         self.sources_size = sources_size
@@ -125,9 +143,35 @@ def forward(self, model_params, batch):
         stream_cell_tokens = checkpoint(
             self.embed_engine, batch, model_params.pe_embed, use_reentrant=False
         )
+        cell_lens = torch.sum(batch.tokens_lens, 2).flatten()
+        batch_num_cells = batch.tokens_lens.shape[-1]
+        if batch_num_cells == self.local_num_healpix_cells:
+            # The data pipeline already constructed only this rank's HEALPix
+            # cells, so its packed tokens are local without another selection.
+            local_cell_lens = cell_lens
+        elif batch_num_cells == self.num_healpix_cells:
+            # Backward-compatible path for batches constructed with the full
+            # global grid.
+            stream_cell_tokens, local_cell_lens = select_packed_cell_shard(
+                stream_cell_tokens,
+                cell_lens,
+                self.num_healpix_cells,
+                self.local_cell_start,
+                self.local_cell_end,
+            )
+        else:
+            raise ValueError(
+                f"batch has {batch_num_cells} HEALPix cells; expected either "
+                f"{self.local_num_healpix_cells} local or {self.num_healpix_cells} global cells"
+            )
 
         tokens_global, posteriors = checkpoint(
-            self.assimilate_local, model_params, stream_cell_tokens, batch, use_reentrant=False
+            self.assimilate_local,
+            model_params,
+            stream_cell_tokens,
+            batch,
+            local_cell_lens,
+            use_reentrant=False,
         )
 
         tokens_global = checkpoint(
@@ -153,7 +197,14 @@ def interpolate_latents(self, tokens: torch.Tensor) -> (torch.Tensor, torch.Tens
 
         return tokens, posteriors
 
-    def assimilate_local_project_chunked(self, tokens, tokens_global, cell_lens, q_cells_lens):
+    def assimilate_local_project_chunked(
+        self,
+        tokens,
+        tokens_global,
+        cell_lens,
+        q_cells_lens,
+        num_cells_per_sample=None,
+    ):
         """
         Apply the local assimilation engine and then the
         local-to-global adapter using a chunking in the number of tokens
@@ -163,14 +214,21 @@ def assimilate_local_project_chunked(self, tokens, tokens_global, cell_lens, q_c
         # combined cell lens for all tokens in batch across all input steps
         zero_pad = torch.zeros(1, device=tokens.device, dtype=torch.int32)
 
-        # subdivision factor for required splitting
-        clen = self.num_healpix_cells // (2 if self.cf.healpix_level <= 5 else 8)
+        # Spatial parallelism already reduces a rank to at most 1/8 of the
+        # HEALPix cells. Process that shard in one call so every spatial rank
+        # enters the local FSDP modules the same number of times.
+        num_cells_per_sample = num_cells_per_sample or self.num_healpix_cells
+        if self.spatial_parallel_size > 1:
+            clen = num_cells_per_sample
+        else:
+            clen = self.num_healpix_cells // (2 if self.cf.healpix_level <= 5 else 8)
         tokens_global_unmasked = []
         posteriors = []
+        empty_chunk_dependency = tokens_global.new_zeros(())
 
-        for i in range(cell_lens.shape[0] // clen):
-            # make sure we properly catch all elements in last chunk
-            i_end = (i + 1) * clen if i < (cell_lens.shape[0] // clen) - 1 else cell_lens.shape[0]
+        num_chunks = (cell_lens.shape[0] + clen - 1) // clen
+        for i in range(num_chunks):
+            i_end = min((i + 1) * clen, cell_lens.shape[0])
             l0, l1 = (
                 (0 if i == 0 else cell_lens[: i * clen].cumsum(0)[-1]),
                 cell_lens[:i_end].cumsum(0)[-1],
@@ -181,6 +239,20 @@ def assimilate_local_project_chunked(self, tokens, tokens_global, cell_lens, q_c
             # skip processing of the empty chunk in this case
             # Check if this chunk is empty
             if l0 == l1 or toks.shape[0] == 0:
+                # Spatial ranks must enter FSDP-wrapped local modules in the
+                # same order. Run a one-token zero-valued path and retain a
+                # zero dependency so its backward hooks also execute.
+                dummy_lens = torch.tensor([0, 1], device=tokens.device, dtype=torch.int32)
+                dummy_tokens = tokens.new_zeros((1, tokens.shape[-1]))
+                dummy_tokens = self.ae_local_engine(dummy_tokens, dummy_lens, use_reentrant=False)
+                dummy_tokens, _ = self.interpolate_latents(dummy_tokens)
+                dummy_global = self.ae_local_global_engine(
+                    dummy_tokens,
+                    tokens_global[i * clen : i * clen + 1],
+                    dummy_lens,
+                    dummy_lens,
+                )
+                empty_chunk_dependency = empty_chunk_dependency + dummy_global.sum() * 0
                 continue
 
             toks_global = tokens_global[i * clen : i_end]
@@ -210,10 +282,13 @@ def assimilate_local_project_chunked(self, tokens, tokens_global, cell_lens, q_c
             tokens_global_unmasked += [toks_global_unmasked]
 
         if len(tokens_global_unmasked) == 0:
-            assert False, "Not yet implemented"
-        tokens_global_unmasked = torch.cat(tokens_global_unmasked)
+            tokens_global_unmasked = tokens_global.new_empty(
+                (0, tokens_global.shape[-2], tokens_global.shape[-1])
+            )
+        else:
+            tokens_global_unmasked = torch.cat(tokens_global_unmasked)
 
-        return tokens_global_unmasked, posteriors
+        return tokens_global_unmasked, posteriors, empty_chunk_dependency
 
     def aggregation_engine_unmasked(
         self,
@@ -273,7 +348,11 @@ def aggregation_engine_unmasked(
         return tokens_global_unmasked
 
     def assimilate_local(
-        self, model_params, tokens: torch.Tensor, batch: ModelBatch
+        self,
+        model_params,
+        tokens: torch.Tensor,
+        batch: ModelBatch,
+        cell_lens_local: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """
         Processes embedded tokens locally and prepares them for the global assimilation
@@ -287,7 +366,23 @@ def assimilate_local(
             Tokens for global assimilation
         """
 
-        cell_lens = torch.sum(batch.tokens_lens, 2).flatten()
+        tokens_lens_global = batch.tokens_lens
+        batch_num_cells = tokens_lens_global.shape[-1]
+        if batch_num_cells == self.local_num_healpix_cells:
+            if self.spatial_parallel_size > 1:
+                tokens_lens_global = torch.cat(
+                    all_gather(
+                        tokens_lens_global,
+                        group=self.spatial_parallel_group,
+                    ),
+                    dim=-1,
+                )
+        elif batch_num_cells != self.num_healpix_cells:
+            raise ValueError(
+                f"batch has {batch_num_cells} HEALPix cells; expected either "
+                f"{self.local_num_healpix_cells} local or {self.num_healpix_cells} global cells"
+            )
+        cell_lens = torch.sum(tokens_lens_global, 2).flatten()
 
         num_steps_input = batch.get_num_source_steps()
         rs = num_steps_input * len(batch)
@@ -297,34 +392,75 @@ def assimilate_local(
         pos_enc = positional_encoding_harmonic
         tokens_global_register_class = pos_enc(self.q_cells.repeat(rs, num_extra_tokens, 1))
 
+        # Direct calls retain the old API and perform the shard selection here.
+        # ``forward`` passes an already-sharded stream_cell_tokens tensor.
+        if cell_lens_local is None:
+            if batch_num_cells == self.local_num_healpix_cells:
+                cell_lens_local = torch.sum(batch.tokens_lens, 2).flatten()
+            else:
+                tokens, cell_lens_local = select_packed_cell_shard(
+                    tokens,
+                    cell_lens,
+                    self.num_healpix_cells,
+                    self.local_cell_start,
+                    self.local_cell_end,
+                )
+
+        pe_global_local = model_params.pe_global[self.local_cell_start : self.local_cell_end]
+
         # TODO: re-enable or remove ae_local_queries_per_cell
         if self.cf.ae_local_queries_per_cell:
-            tokens_global = (self.q_cells + model_params.pe_global).repeat(rs, 1, 1)
+            q_cells_local = self.q_cells[self.local_cell_start : self.local_cell_end]
+            tokens_global = (q_cells_local + pe_global_local).repeat(rs, 1, 1)
         else:
-            num_tokens = self.num_healpix_cells
-            tokens_global = self.q_cells.repeat(num_tokens, 1, 1) + model_params.pe_global
+            tokens_global = (
+                self.q_cells.repeat(self.local_num_healpix_cells, 1, 1) + pe_global_local
+            )
             tokens_global = tokens_global.repeat(rs, 1, 1)
 
         # apply local assimilation engine and project onto global latent vectors
-        tokens_global_unmasked, posteriors = self.assimilate_local_project_chunked(
-            tokens, tokens_global, cell_lens, model_params.q_cells_lens
+        tokens_global_unmasked, posteriors, empty_chunk_dependency = (
+            self.assimilate_local_project_chunked(
+                tokens,
+                tokens_global,
+                cell_lens_local,
+                model_params.q_cells_lens,
+                self.local_num_healpix_cells,
+            )
+        )
+        tokens_global = tokens_global + empty_chunk_dependency
+
+        # Restore a dense local cell tensor before gathering. This gives every
+        # rank the same gather shape and preserves the full autograd graph.
+        local_mask = cell_lens_local.to(torch.bool)
+        tokens_global[local_mask] = tokens_global_unmasked.to(tokens_global.dtype)
+        tokens_global = tokens_global.reshape(
+            rs,
+            self.local_num_healpix_cells,
+            self.q_cells.shape[-2],
+            self.q_cells.shape[-1],
         )
+        if self.spatial_parallel_size > 1:
+            tokens_global = torch.cat(
+                all_gather(tokens_global, group=self.spatial_parallel_group),
+                dim=1,
+            )
+
+        # Recover packed, globally ordered non-empty cells for query aggregation.
+        cell_mask = cell_lens.reshape(rs, self.num_healpix_cells).to(torch.bool)
+        tokens_global_unmasked = tokens_global[cell_mask]
 
         # apply aggregation engine on unmasked tokens
         tokens_global_unmasked = self.aggregation_engine_unmasked(
             tokens_global_unmasked,
             tokens_global_register_class,
-            batch.tokens_lens,
+            tokens_lens_global,
             rope_cell_coords=model_params.rope_cell_coords,
         )
 
         # final processing
 
-        tokens_global = (
-            torch.permute(tokens_global, [1, 0, 2])
-            .squeeze()
-            .reshape(rs, self.num_healpix_cells, -1)
-        )
+        tokens_global = tokens_global.reshape(rs, self.num_healpix_cells, -1)
         # TODO, TODO, TODO: do we need this
         tokens_global = torch.cat([tokens_global_register_class, tokens_global], dim=1)
 
diff --git a/src/weathergen/model/engines.py b/src/weathergen/model/engines.py
index fde31213b6..03669c7423 100644
--- a/src/weathergen/model/engines.py
+++ b/src/weathergen/model/engines.py
@@ -115,7 +115,12 @@ def forward(self, batch, pe_embed):
         )
         " Increase ae_local_max_tokens_per_cell in config."
 
-        if batch.tokens_lens.shape[2] == 1:
+        if not x_embeds:
+            # A spatial rank can legitimately own a domain with no observations
+            # for this sample. Keep an empty tensor so all ranks can continue to
+            # the synchronized local-assimilation path.
+            return tokens_all
+        elif batch.tokens_lens.shape[2] == 1:
             # trivial with one stream
             tokens_all = torch.cat(x_embeds)
 
diff --git a/src/weathergen/model/spatial_parallel.py b/src/weathergen/model/spatial_parallel.py
new file mode 100644
index 0000000000..893700b17a
--- /dev/null
+++ b/src/weathergen/model/spatial_parallel.py
@@ -0,0 +1,47 @@
+# (C) Copyright 2025 WeatherGenerator contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+import torch
+
+
+def select_packed_cell_shard(
+    tokens: torch.Tensor,
+    cell_lens: torch.Tensor,
+    num_cells: int,
+    cell_start: int,
+    cell_end: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Select a HEALPix-cell range from a cell-packed token tensor.
+
+    ``cell_lens`` is flattened in ``(input_step, sample, cell)`` order. Token
+    counts vary by cell, so slicing the first tensor dimension directly would
+    split cells. The returned tokens remain packed in the same order, restricted
+    to ``[cell_start, cell_end)`` for every input-step/sample row.
+    """
+
+    if cell_lens.numel() % num_cells:
+        raise ValueError("cell_lens does not contain a whole number of HEALPix grids")
+    if not 0 <= cell_start < cell_end <= num_cells:
+        raise ValueError(
+            f"invalid HEALPix cell range [{cell_start}, {cell_end}) for {num_cells} cells"
+        )
+
+    cell_lens_2d = cell_lens.reshape(-1, num_cells)
+    selected_cells = torch.zeros_like(cell_lens_2d, dtype=torch.bool)
+    selected_cells[:, cell_start:cell_end] = True
+    selected_tokens = torch.repeat_interleave(
+        selected_cells.flatten(), cell_lens.to(dtype=torch.long)
+    )
+    if selected_tokens.numel() != tokens.shape[0]:
+        raise ValueError(
+            f"packed token length mismatch: cell_lens describes {selected_tokens.numel()} "
+            f"tokens, tensor has {tokens.shape[0]}"
+        )
+
+    return tokens[selected_tokens], cell_lens_2d[:, cell_start:cell_end].flatten()
diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py
index 276da0bd67..e46c883f09 100644
--- a/src/weathergen/train/trainer.py
+++ b/src/weathergen/train/trainer.py
@@ -47,7 +47,7 @@
     get_batch_size_from_config,
     get_target_idxs_from_cfg,
 )
-from weathergen.utils.distributed import is_root
+from weathergen.utils.distributed import get_encoder_spatial_parallel_size, is_root
 from weathergen.utils.performance import NullThroughputTracker, ThroughputTracker, nvtx_range
 from weathergen.utils.train_logger import TrainLogger, prepare_losses_for_logging
 from weathergen.utils.utils import get_dtype
@@ -95,7 +95,7 @@ def get_batch_size_total(self, batch_size_per_gpu) -> int:
         """
         Get total, effective batch size across all DDP ranks
         """
-        return self.world_size_original * batch_size_per_gpu
+        return self.data_parallel_world_size_original * batch_size_per_gpu
 
     def init(self, cf: Config, devices):
         # pylint: disable=attribute-defined-outside-init
@@ -151,6 +151,19 @@ def init(self, cf: Config, devices):
         # world_size gets overwritten by current setting during init_ddp()
         self.world_size_original = cf.get("world_size_original", cf.get("world_size", None))
         cf.world_size_original = self.world_size_original
+        spatial_parallel_size = get_encoder_spatial_parallel_size(cf)
+        cf.data_parallel_world_size = cf.world_size // spatial_parallel_size
+        spatial_parallel_size_original = cf.get(
+            "encoder_spatial_parallel_size_original", spatial_parallel_size
+        )
+        if self.world_size_original % spatial_parallel_size_original:
+            raise ValueError(
+                "world_size_original must be divisible by encoder_spatial_parallel_size_original"
+            )
+        self.data_parallel_world_size_original = (
+            self.world_size_original // spatial_parallel_size_original
+        )
+        cf.encoder_spatial_parallel_size_original = spatial_parallel_size_original
 
         self.log_grad_norms = cf.train_logging.get("log_grad_norms", False)
 
@@ -349,7 +362,7 @@ def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None):
         self.lr_scheduler = LearningRateScheduler(
             self.optimizer,
             self.batch_size_per_gpu,
-            cf.world_size,
+            cf.data_parallel_world_size,
             cf.general.istep,
             lr_steps,
             self.training_cfg.learning_rate_scheduling,
@@ -368,13 +381,17 @@ def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None):
             mini_epoch_base = int(self.cf.general.istep / len(self.data_loader))
         else:
             len_per_rank = (
-                max(1, len(self.dataset) // (self.world_size_original * self.batch_size_per_gpu))
+                max(
+                    1,
+                    len(self.dataset)
+                    // (self.data_parallel_world_size_original * self.batch_size_per_gpu),
+                )
             ) * self.batch_size_per_gpu
             mini_epoch_base = int(
                 self.cf.general.istep
                 / (
                     min(len_per_rank, self.training_cfg.samples_per_mini_epoch)
-                    * self.world_size_original
+                    * self.data_parallel_world_size_original
                 )
             )
 
diff --git a/src/weathergen/utils/distributed.py b/src/weathergen/utils/distributed.py
index af467a3a94..77b48a54d5 100644
--- a/src/weathergen/utils/distributed.py
+++ b/src/weathergen/utils/distributed.py
@@ -12,6 +12,7 @@
 import torch.distributed as dist
 
 SYNC_TIMEOUT_SEC = 60 * 60  # 1 hour
+_ENCODER_SPATIAL_GROUPS: dict[int, tuple[dist.ProcessGroup, int]] = {}
 
 
 def is_root(pg: dist.ProcessGroup | None = None) -> bool:
@@ -59,6 +60,57 @@ def get_rank() -> int:
     return dist.get_rank()
 
 
+def get_encoder_spatial_parallel_size(cf) -> int:
+    """Return and validate the configured encoder spatial-parallel size."""
+
+    size = int(cf.get("encoder_spatial_parallel_size", 1))
+    if size < 1:
+        raise ValueError("encoder_spatial_parallel_size must be at least 1")
+
+    world_size = get_world_size()
+    if size > world_size:
+        raise ValueError(
+            f"encoder_spatial_parallel_size ({size}) exceeds world_size ({world_size})"
+        )
+    if world_size % size:
+        raise ValueError(
+            f"world_size ({world_size}) must be divisible by encoder_spatial_parallel_size ({size})"
+        )
+    return size
+
+
+def get_encoder_spatial_parallel_group(cf) -> tuple[dist.ProcessGroup | None, int]:
+    """Create the consecutive-rank process groups used to shard HEALPix cells.
+
+    All ranks call ``new_group`` in the same order. The returned rank is local to
+    the spatial group. A size of one deliberately avoids creating a process group.
+    """
+
+    size = get_encoder_spatial_parallel_size(cf)
+    if size == 1:
+        return None, 0
+    if not _is_distributed_initialized():
+        raise RuntimeError("encoder spatial parallelism requires torch.distributed")
+
+    cached = _ENCODER_SPATIAL_GROUPS.get(size)
+    if cached is not None:
+        return cached
+
+    world_size = dist.get_world_size()
+    global_rank = dist.get_rank()
+    own_group = None
+    for first_rank in range(0, world_size, size):
+        ranks = list(range(first_rank, first_rank + size))
+        group = dist.new_group(ranks=ranks)
+        if global_rank in ranks:
+            own_group = group
+
+    assert own_group is not None
+    result = (own_group, global_rank % size)
+    _ENCODER_SPATIAL_GROUPS[size] = result
+    return result
+
+
 def ddp_average(data: torch.Tensor) -> torch.Tensor:
     """
     Average a tensor across DDP ranks
diff --git a/tests/test_encoder_spatial_parallel.py b/tests/test_encoder_spatial_parallel.py
new file mode 100644
index 0000000000..923245c575
--- /dev/null
+++ b/tests/test_encoder_spatial_parallel.py
@@ -0,0 +1,146 @@
+# (C) Copyright 2025 WeatherGenerator contributors.
+#
+# This software is licensed under the terms of the Apache Licence Version 2.0
+# which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# In applying this licence, ECMWF does not waive the privileges and immunities
+# granted to it by virtue of its status as an intergovernmental organisation
+# nor does it submit to any jurisdiction.
+
+import numpy as np
+import pytest
+import torch
+
+from weathergen.datasets.healpix_domain import build_local_healpix_cell_splits
+from weathergen.model.spatial_parallel import select_packed_cell_shard
+from weathergen.utils import distributed
+
+
+def test_local_healpix_construction_matches_global_cell_slices():
+    num_cells = 48
+    cell_ids = np.repeat(np.arange(num_cells), np.arange(num_cells) % 3 + 1)
+    rng = np.random.default_rng(7)
+    cell_ids = cell_ids[rng.permutation(len(cell_ids))]
+    global_cells = build_local_healpix_cell_splits(
+        cell_ids,
+        num_cells,
+        cell_start=0,
+        cell_end=num_cells,
+    )
+    cells_per_rank = len(global_cells) // 4
+
+    local_cells_all = []
+    for spatial_rank in range(4):
+        cell_start = spatial_rank * cells_per_rank
+        cell_end = cell_start + cells_per_rank
+        local_cells = build_local_healpix_cell_splits(
+            cell_ids,
+            num_cells,
+            cell_start=cell_start,
+            cell_end=cell_end,
+        )
+
+        assert len(local_cells) == cells_per_rank
+        for local_cell, global_cell in zip(
+            local_cells,
+            global_cells[cell_start:cell_end],
+            strict=True,
+        ):
+            np.testing.assert_array_equal(local_cell, global_cell)
+        local_cells_all.extend(local_cells)
+
+    assert len(local_cells_all) == len(global_cells)
+    for local_cell, global_cell in zip(local_cells_all, global_cells, strict=True):
+        np.testing.assert_array_equal(local_cell, global_cell)
+
+
+def test_local_healpix_construction_rejects_invalid_range():
+    with pytest.raises(ValueError, match="invalid HEALPix cell range"):
+        build_local_healpix_cell_splits(
+            np.arange(12),
+            num_cells=12,
+            cell_start=6,
+            cell_end=13,
+        )
+
+
+def test_local_healpix_construction_handles_empty_domain():
+    cell_splits = build_local_healpix_cell_splits(
+        np.array([0, 1, 2], dtype=np.int64),
+        num_cells=12,
+        cell_start=6,
+        cell_end=9,
+    )
+
+    assert len(cell_splits) == 3
+    assert all(cell.dtype == np.int64 and cell.size == 0 for cell in cell_splits)
+
+
+def test_select_packed_cell_shard_preserves_cell_boundaries_across_rows():
+    cell_lens = torch.tensor(
+        [
+            [1, 0, 2, 1, 3, 0, 1, 2],
+            [0, 2, 1, 0, 1, 2, 0, 1],
+        ],
+        dtype=torch.int32,
+    )
+    tokens = torch.arange(cell_lens.sum(), dtype=torch.float32).unsqueeze(1)
+
+    shard, shard_lens = select_packed_cell_shard(
+        tokens, cell_lens.flatten(), num_cells=8, cell_start=2, cell_end=4
+    )
+
+    assert shard_lens.tolist() == [2, 1, 1, 0]
+    assert shard.squeeze(1).tolist() == [1, 2, 3, 12]
+
+
+def test_eight_shards_cover_every_packed_token_once_and_keep_gradients():
+    num_cells = 16
+    cell_lens = torch.tensor(
+        [
+            [0, 1, 2, 0, 1, 3, 0, 2, 1, 0, 2, 1, 0, 1, 2, 1],
+            [1, 0, 1, 2, 0, 1, 2, 0, 3, 1, 0, 1, 2, 0, 1, 1],
+        ],
+        dtype=torch.int32,
+    )
+    tokens = torch.arange(cell_lens.sum(), dtype=torch.float32, requires_grad=True)
+    shard_width = num_cells // 8
+
+    selected = []
+    for rank in range(8):
+        shard, _ = select_packed_cell_shard(
+            tokens,
+            cell_lens.flatten(),
+            num_cells,
+            rank * shard_width,
+            (rank + 1) * shard_width,
+        )
+        selected.append(shard)
+        shard.sum().backward(retain_graph=rank < 7)
+
+    assert sum(shard.numel() for shard in selected) == tokens.numel()
+    assert torch.equal(tokens.grad, torch.ones_like(tokens))
+
+
+@pytest.mark.parametrize(
+    ("num_cells", "cell_start", "cell_end"),
+    [(8, -1, 1), (8, 3, 3), (8, 0, 9)],
+)
+def test_select_packed_cell_shard_rejects_invalid_ranges(num_cells, cell_start, cell_end):
+    with pytest.raises(ValueError, match="invalid HEALPix cell range"):
+        select_packed_cell_shard(
+            torch.arange(num_cells),
+            torch.ones(num_cells, dtype=torch.int32),
+            num_cells,
+            cell_start,
+            cell_end,
+        )
+
+
+def test_spatial_parallel_size_requires_whole_rank_groups(monkeypatch):
+    monkeypatch.setattr(distributed, "get_world_size", lambda: 16)
+    assert distributed.get_encoder_spatial_parallel_size({"encoder_spatial_parallel_size": 4}) == 4
+    assert distributed.get_encoder_spatial_parallel_size({"encoder_spatial_parallel_size": 8}) == 8
+
+    with pytest.raises(ValueError, match="must be divisible"):
+        distributed.get_encoder_spatial_parallel_size({"encoder_spatial_parallel_size": 6})