Deepseek Engram optimization. by LiJunscs · Pull Request #1147 · flagos-ai/FlagScale

LiJunscs · 2026-03-12T11:27:36Z

PR Category

[Train]

PR Types

[Improvements]

PR Description

AlltoAll communication when compute multi_head_embedding.
Precompute multi_head_embedding.
Optional offloading embedding's optimizer states.

CLAassistant · 2026-03-12T12:59:24Z

All committers have signed the CLA.

Copilot

Pull request overview

This PR improves DeepSeek Engram training performance by adding embedding-parallel communication options (all-to-all / all-reduce / offload), enabling precomputation/overlap of multi-head embedding work, and wiring new Engram-parallel arguments through initialization and example configs.

Changes:

Add Engram embedding parallel method/size arguments and pass them into distributed initialization.
Introduce an embedding-parallel EngramMemory implementation and integrate it into Engram’s multi-head embedding + sharded checkpointing.
Add embedding precompute/caching (and attempt overlap across layers) plus new DeepSeek Engram Hydra configs.

Reviewed changes

Copilot reviewed 11 out of 11 changed files in this pull request and generated 12 comments.

Show a summary per file

File	Description
tests/functional_tests/train/deepseek/gold_values/tp2_pp2_ep2_engram.json	Reformat/update stored gold loss JSON.
flagscale/train/megatron/training/initialize.py	Plumb `engram_embedding_parallel_size` into model-parallel initialization.
flagscale/train/megatron/training/arguments_fs.py	Add Engram embedding parallel CLI args + validation/warnings.
flagscale/train/megatron/train_engram.py	Switch to `parallel_state` and add an extra TP token broadcast path.
flagscale/models/megatron/engram/multi_head_embedding.py	Add `EngramMemory` and route MultiHeadEmbedding through it; add sharded state dict.
flagscale/models/megatron/engram/engram_transformer_layer.py	Add embedding precompute hook and enable sharded state dict.
flagscale/models/megatron/engram/engram_model.py	Implement `build_schedule_plan` for Engram model.
flagscale/models/megatron/engram/engram_config.py	Add Engram embedding parallel config fields.
flagscale/models/megatron/engram/engram.py	Add embedding precompute/cache + sharded state dict plumbing; rename embedding member to `memory`.
examples/deepseek_v3/conf/train_engram.yaml	New top-level Hydra entry for Engram training run.
examples/deepseek_v3/conf/train/engram.yaml	New DeepSeek Engram training preset including embedding-parallel settings.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

+    if not (parallel_state.get_pipeline_model_parallel_world_size == 1 or parallel_state.is_pipeline_first_stage()):
+        if parallel_state.get_tensor_model_parallel_rank() == 0:
+            torch.distributed.broadcast(batch["tokens"], src=parallel_state.get_tensor_model_parallel_src_rank(), group=parallel_state.get_tensor_model_parallel_group())
+        else:
+            tokens = torch.empty_like(batch["labels"])
+            torch.distributed.broadcast(tokens, src=parallel_state.get_tensor_model_parallel_src_rank(), group=parallel_state.get_tensor_model_parallel_group())
+            batch["tokens"] = tokens


 from megatron.core.utils import get_attr_wrapped_model, StragglerDetector
 from megatron.core.tokenizers.text.utils.build_tokenizer import build_tokenizer
-from megatron.core import mpu
+from megatron.core import parallel_state


+        (self.vocab_start_index, self.vocab_end_index) = (
+            VocabUtility.vocab_range_from_global_vocab_size(
+                self.num_embeddings, get_pg_rank(self.embedding_parallel_group), get_pg_size(self.embedding_parallel_group)
+            )


+                    rank=get_pg_rank(self.tp_group),
+                    world_size=get_pg_size(self.tp_group),


+        input_ids = input_ids.view(-1)
+        routing_map = input_ids // self.num_embeddings_per_partition
+        # [num_partitions], number of tokens assigned to each partition from the current rank's input.
+        num_tokens_per_partition = torch.histc(routing_map, bins=self.embedding_parallel_size, min=0, max=self.embedding_parallel_size)


+                    warnings.warn(f"[rank0]: We do not recomend using allreduce for engram embedding, this is deprecated and will be removed in later version.", DeprecationWarning)
+                if self.args.engram_embedding_parallel_size is not None:


+        if parallel_state.get_tensor_model_parallel_rank() == 0:
+            torch.distributed.broadcast(batch["tokens"], src=parallel_state.get_tensor_model_parallel_src_rank(), group=parallel_state.get_tensor_model_parallel_group())
+        else:
+            tokens = torch.empty_like(batch["labels"])


+        self,
+        prefix: str = '',
+        sharded_offsets: Tuple[Tuple[int, int, int]] = (),
+        metadata: Optional[dict] = None,** kwargs,


+            assert not self.args.use_megatron_fsdp, "Megatron FSDP does not be supported yet, looking forward to later version."
+            assert not self.args.init_model_with_meta_device, "Init_model_with_meta_device does not be supported yet, looking forward to later version."


+    backend: torchrun
+    nnodes: 3
+    nproc_per_node: 8
+    hostfile: hostfile # Select an available hostfile. Like ip_1 slosts=8\nip_2 slost=8...


lxd-cumt

LGTM

1. All2All when compute embeddings. 2. Precompute multi_head_embedding.

…timizer.

…l and offload to be enabled simultaneously.

1. fix: fix bug of engram memory sharded_state_dict when layers are same. 2. feat: add new golden value of engram when alltoall.

Copilot

Pull request overview

This PR improves DeepSeek Engram training by introducing an all-to-all–based embedding parallel path, adding embedding precompute overlap, and wiring in optional optimizer-state offloading and sharded-checkpoint support for Engram components.

Changes:

Add Engram embedding parallel configuration (size/method) and optional embedding optimizer-state offloading flags.
Implement an Engram all-to-all embedding path (EngramMemory) plus embedding precompute/caching to overlap compute.
Update parameter-norm logic and distributed initialization to account for Engram embedding parallel groups; refresh functional test config/gold values and add example configs.

Reviewed changes

Copilot reviewed 12 out of 12 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
tests/functional_tests/train/deepseek/gold_values/tp2_pp2_ep2_engram.json	Update expected loss gold values for the new Engram behavior/config.
tests/functional_tests/train/deepseek/conf/train/tp2_pp2_ep2_engram.yaml	Add Engram embedding parallel settings and adjust attention/MoE knobs for the functional test.
flagscale/train/megatron/training/utils.py	Track Engram embedding params separately and reduce their norm over an Engram-specific MP group.
flagscale/train/megatron/training/initialize.py	Plumb `engram_embedding_parallel_size` into model-parallel initialization.
flagscale/train/megatron/training/arguments_fs.py	Add Engram CLI args and post-validate logic for embedding parallel method/size and offload constraints.
flagscale/models/megatron/engram/multi_head_embedding.py	Implement `EngramMemory` all-to-all embedding path and route `MultiHeadEmbedding` between allreduce/alltoall methods; add sharded-state-dict handling.
flagscale/models/megatron/engram/engram_transformer_layer.py	Add embedding precompute hook and enable sharded-state-dict passthrough with a non-homogeneous layers flag.
flagscale/models/megatron/engram/engram_model.py	Add async hash computation wrapper and extend schedule-plan building to pass lazy hash inputs.
flagscale/models/megatron/engram/engram_config.py	Extend `EngramConfig` with embedding-parallel and offloading options.
flagscale/models/megatron/engram/engram.py	Add embedding precompute/caching via CUDA stream/event and implement sharded-state-dict composition.
examples/deepseek_v3/conf/train_engram.yaml	New runnable example wrapper config for Engram training.
examples/deepseek_v3/conf/train/engram.yaml	New full Engram training config (DeepSeek v3) including embedding-parallel settings.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-04-14T03:17:03Z

+                warnings.warn("Offloading embedding optimizer states will offload all embedding optimizer states to CPU, which may cause slowdown. " \
+                "Please make sure this is what you want. This is typically used to save GPU memory when Engram embedding is large while accelerators are limited." \
+                "If you do not want to offload all embedding optimizer states to CPU, please disable this and set the --optimizer-offload-fraction to a value less than 1 to offload part of the optimizer states to CPU." \
+                "Of cource you can set the --optimizer-offload-fraction to offload other params meanwhile enable this to offload all embedding optimizer states to CPU.")                


Typo in warning text: "Of cource" should be "Of course".

Suggested change

"Of cource you can set the --optimizer-offload-fraction to offload other params meanwhile enable this to offload all embedding optimizer states to CPU.")

"Of course you can set the --optimizer-offload-fraction to offload other params meanwhile enable this to offload all embedding optimizer states to CPU.")

Copilot · 2026-04-14T03:17:03Z

+        self,
+        prefix: str = '',
+        sharded_offsets: Tuple[Tuple[int, int, int]] = (),
+        metadata: Optional[dict] = None,** kwargs,


EngramMemory.sharded_state_dict has invalid Python syntax in the parameter list: metadata: Optional[dict] = None,** kwargs, (space in ** kwargs and missing space after comma). This will raise a SyntaxError and prevent the module from importing. Change the signature to a valid **kwargs form (e.g., metadata: Optional[dict] = None, **kwargs).

Suggested change

metadata: Optional[dict] = None,** kwargs,

metadata: Optional[dict] = None, **kwargs,

Copilot · 2026-04-14T03:17:04Z

+        # This is usefule when all layer are same, the TransformerBlock will be homogeneous, it generate sharded_state_dict will same keys for all layer and need all layers have the same structure.
+        # The layer has engram module does not fit this assumption.
+        # If the flag is set to True, the sharded_state_dict will use layer_number to generate different keys for different layer, which is same to models has dense layer leading and moe layer following.
+        # Actually, engram really causes the layers to be non-homogeneous.


sharded_state_dict() writes into metadata["non_homogeneous_layers"] without checking whether metadata is None. Since the method signature allows metadata: dict | None = None, this will raise a TypeError if callers omit metadata. Ensure metadata is initialized (e.g., set metadata = {} if metadata is None else metadata) before mutation.

Suggested change

# Actually, engram really causes the layers to be non-homogeneous.

# Actually, engram really causes the layers to be non-homogeneous.

metadata = {} if metadata is None else metadata

Actually, metadata will never be None.

Copilot · 2026-04-14T03:17:04Z

+    def pre_compute_embedding(self, input_ids: torch.Tensor):
+        """
+        Pre-compute the multi-head embedding for the given input IDs.
+        This can be called before the forward pass to warm up the embedding cache.
+        """
+        assert input_ids is not None, "Input ids can not be None for EngramModel"
+        self.embedding_stream.synchronize()  # Ensure previous computations on the stream are finished
+        with torch.cuda.stream(self.embedding_stream):
+            embedding_result = self.memory(input_ids).flatten(start_dim=-2)
+        embedding_event = torch.cuda.Event()
+        embedding_event.record(self.embedding_stream)
+        self.embedding_cache = (embedding_result, embedding_event)


pre_compute_embedding() unconditionally uses CUDA APIs (self.embedding_stream.synchronize(), torch.cuda.stream(...), torch.cuda.Event()) but self.embedding_stream is only created when torch.cuda.is_available(). If Engram is instantiated in a CPU-only run (or CUDA is disabled), this will raise an exception. Guard this method (and/or the call sites) so it either becomes a no-op on CPU or falls back to synchronous computation without CUDA streams/events.

Copilot · 2026-04-14T03:17:04Z

+                    warnings.warn(
+                        "[rank0]: If set the embedding_parallel_method to allreduce, " \
+                        "the embedding module will be the tensor_parallel.layers.VocabParallelEmbedding with tensor_parallel." \
+                        "So the embedding_parallel_size is useless and set to None."
+                    )


This warning message is labeled with "[rank0]" but is emitted on every rank (it isn't guarded by rank == 0). This can spam logs and is misleading. Either guard this warning with the same rank == 0 check, or remove the "[rank0]" prefix from the message.

Suggested change

warnings.warn(

"[rank0]: If set the embedding_parallel_method to allreduce, " \

"the embedding module will be the tensor_parallel.layers.VocabParallelEmbedding with tensor_parallel." \

"So the embedding_parallel_size is useless and set to None."

)

if self.args.rank == 0:

warnings.warn(

"[rank0]: If set the embedding_parallel_method to allreduce, " \

"the embedding module will be the tensor_parallel.layers.VocabParallelEmbedding with tensor_parallel." \

"So the embedding_parallel_size is useless and set to None."

)

Copilot · 2026-04-14T03:17:04Z

+
+
+class EngramMemory(nn.Module):
+    """Embedding parallelized in the vocabulay dimension.


Typo in docstring: "vocabulay" should be "vocabulary".

Suggested change

"""Embedding parallelized in the vocabulay dimension.

"""Embedding parallelized in the vocabulary dimension.

Copilot · 2026-04-14T03:17:05Z

+        # Engram let the layers be non-homogeneous, so we need to set the flag in metadata to let the sharded state dict logic know.
+        # This is usefule when all layer are same, the TransformerBlock will be homogeneous, it generate sharded_state_dict will same keys for all layer and need all layers have the same structure.
+        # The layer has engram module does not fit this assumption.
+        # If the flag is set to True, the sharded_state_dict will use layer_number to generate different keys for different layer, which is same to models has dense layer leading and moe layer following.
+        # Actually, engram really causes the layers to be non-homogeneous.


Typo in comment: "usefule" should be "useful".

Suggested change

# Engram let the layers be non-homogeneous, so we need to set the flag in metadata to let the sharded state dict logic know.

# This is usefule when all layer are same, the TransformerBlock will be homogeneous, it generate sharded_state_dict will same keys for all layer and need all layers have the same structure.

# The layer has engram module does not fit this assumption.

# If the flag is set to True, the sharded_state_dict will use layer_number to generate different keys for different layer, which is same to models has dense layer leading and moe layer following.

# Actually, engram really causes the layers to be non-homogeneous.

# Engram can make the layers non-homogeneous, so we need to set a flag in

# metadata to let the sharded state dict logic know.

# This is useful because when all layers are the same, TransformerBlock is

# homogeneous and generates the same sharded_state_dict keys for all layers,

# which requires all layers to have the same structure.

# Layers with an Engram module do not fit this assumption.

# If this flag is set to True, sharded_state_dict will use layer_number to

# generate different keys for different layers, similar to models with dense

# layers first and MoE layers following.

Copilot

Pull request overview

Copilot reviewed 12 out of 12 changed files in this pull request and generated 4 comments.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

Copilot · 2026-04-14T14:22:17Z

+        self.embedding_stream.synchronize()  # Ensure previous computations on the stream are finished
+        with torch.cuda.stream(self.embedding_stream):
+            embedding_result = self.memory(input_ids).flatten(start_dim=-2)
+        embedding_event = torch.cuda.Event()
+        embedding_event.record(self.embedding_stream)
+        self.embedding_cache = (embedding_result, embedding_event)
+


pre_compute_embedding() unconditionally uses self.embedding_stream (synchronize() and torch.cuda.stream(...)), but self.embedding_stream is set only when torch.cuda.is_available(). If this method is called in a CPU-only environment (or before CUDA is initialized), it will raise an attribute error / CUDA error. Guard for self.embedding_stream is None (fallback to synchronous compute without streams/events) or assert CUDA availability with a clearer error message.

Suggested change

self.embedding_stream.synchronize() # Ensure previous computations on the stream are finished

with torch.cuda.stream(self.embedding_stream):

embedding_result = self.memory(input_ids).flatten(start_dim=-2)

embedding_event = torch.cuda.Event()

embedding_event.record(self.embedding_stream)

self.embedding_cache = (embedding_result, embedding_event)

embedding_stream = getattr(self, "embedding_stream", None)

if embedding_stream is not None and torch.cuda.is_available():

embedding_stream.synchronize() # Ensure previous computations on the stream are finished

with torch.cuda.stream(embedding_stream):

embedding_result = self.memory(input_ids).flatten(start_dim=-2)

embedding_event = torch.cuda.Event()

embedding_event.record(embedding_stream)

else:

embedding_result = self.memory(input_ids).flatten(start_dim=-2)

embedding_event = None

self.embedding_cache = (embedding_result, embedding_event)

Copilot · 2026-04-14T14:22:18Z

+  engram_seed: 0
+  engram_kernel_size: 4
+  engram_hc_mult: 1
+  engram_embedding_parallel_method: alltoall # alltoall, allreduce, offload


The inline comment lists offload as a valid value for engram_embedding_parallel_method, but the CLI/config choices added in this PR are only alltoall and allreduce (offloading is controlled by a separate boolean flag). Update the comment to avoid misleading users, and optionally document engram_offload_embedding_optimizer_states here instead.

Suggested change

engram_embedding_parallel_method: alltoall # alltoall, allreduce, offload

engram_embedding_parallel_method: alltoall # alltoall, allreduce

# engram_offload_embedding_optimizer_states: true # Optional; controls optimizer-state offloading.

Copilot · 2026-04-14T14:22:18Z

+                    if self.args.engram_embedding_parallel_size is not None:
+                        warnings.warn(
+                            "[rank0]: If set the embedding_parallel_method to allreduce, " \
+                            "the embedding module will be the tensor_parallel.layers.VocabParallelEmbedding with tensor_parallel." \
+                            "So the embedding_parallel_size is useless and set to None."
+                        )
+                        self.args.engram_embedding_parallel_size = None   


engram_embedding_parallel_size is set to None only on rank 0 when engram_embedding_parallel_method == "allreduce". This makes argument values diverge across ranks, which can break any later logic that relies on consistent args (e.g., process-group initialization, padding logic, etc.). Keep the warning on rank 0 if desired, but apply the value normalization (engram_embedding_parallel_size = None) on all ranks (or broadcast a single normalized value).

Suggested change

if self.args.engram_embedding_parallel_size is not None:

warnings.warn(

"[rank0]: If set the embedding_parallel_method to allreduce, " \

"the embedding module will be the tensor_parallel.layers.VocabParallelEmbedding with tensor_parallel." \

"So the embedding_parallel_size is useless and set to None."

)

self.args.engram_embedding_parallel_size = None

if self.args.engram_embedding_parallel_size is not None:

if self.args.rank == 0:

warnings.warn(

"[rank0]: If set the embedding_parallel_method to allreduce, " \

"the embedding module will be the tensor_parallel.layers.VocabParallelEmbedding with tensor_parallel." \

"So the embedding_parallel_size is useless and set to None."

)

self.args.engram_embedding_parallel_size = None

Copilot · 2026-04-14T14:22:18Z

+        if extra_block_kwargs is None:
+            extra_block_kwargs = {
+                "engram_hash_input_ids": engram_hash_input_ids,
+            }


build_schedule_plan() only injects engram_hash_input_ids when extra_block_kwargs is None. If a caller passes a non-empty extra_block_kwargs, the schedule plan will be built without engram_hash_input_ids, but EngramTransformerBlock.forward() requires it. Merge into the provided dict (e.g., copy + setdefault) so engram_hash_input_ids is always present without clobbering user-supplied keys.

Suggested change

if extra_block_kwargs is None:

extra_block_kwargs = {

"engram_hash_input_ids": engram_hash_input_ids,

}

extra_block_kwargs = dict(extra_block_kwargs or {})

extra_block_kwargs.setdefault("engram_hash_input_ids", engram_hash_input_ids)

zhaoyinglia

LGTM

lxd-cumt

LGTM

LiJunscs requested review from aoyulong, heavyrain-lzy and zhaoyinglia as code owners March 12, 2026 11:27

LiJunscs force-pushed the deepseek_related branch 2 times, most recently from e6509b7 to c7d57c3 Compare March 13, 2026 03:17

zhaoyinglia requested review from Copilot and lxd-cumt March 16, 2026 01:45

Copilot started reviewing on behalf of zhaoyinglia March 16, 2026 01:46 View session

Copilot AI reviewed Mar 16, 2026

View reviewed changes

lxd-cumt previously approved these changes Mar 16, 2026

View reviewed changes

LiJunscs dismissed lxd-cumt’s stale review via 1970041 March 16, 2026 12:36

LiJunscs force-pushed the deepseek_related branch 4 times, most recently from 9eded23 to d0a886e Compare April 13, 2026 02:29

LiJunscs requested a review from cyber-pioneer as a code owner April 13, 2026 02:29

LiJunscs force-pushed the deepseek_related branch 2 times, most recently from d060903 to cb45ef4 Compare April 13, 2026 03:31

LiJunscs added 11 commits April 14, 2026 11:08

Optimization of Engram.

63a5dcb

1. All2All when compute embeddings. 2. Precompute multi_head_embedding.

[train] feat: update DeepSeek Engram 27B example yaml

df6f6bb

[train] chore: move engram directory to follow the main

8f528db

[train] chore: delete no use files

9b26315

[train] fix: error path of engram_config

178d7fd

Merge branch 'pr-1107' into optimize_engram

ad494c8

[train] fix: format of some files

6ef717d

[train] chore: update train_engram.yaml example

4bc0f4d

[train] fix: bug of saving and loading sharded_state_dict of Engram

67513a6

[train] chore: remove useless print

66b3167

[train]: fix: fix some corner case of engram.

2a474bf

LiJunscs added 5 commits April 14, 2026 11:08

[train]: fix: fix the bug of engram loss when embedding with a new op…

09fd1aa

…timizer.

[train]: fix: fix bug of log param norm of engram.

480801d

[train]: feat: update the offload logic of engram, allow both paralle…

e603ba7

…l and offload to be enabled simultaneously.

[train]:

c423043

1. fix: fix bug of engram memory sharded_state_dict when layers are same. 2. feat: add new golden value of engram when alltoall.

[train]: feat, update random seed handle of Engram

790ccee

LiJunscs requested a review from Copilot April 14, 2026 03:13

Copilot started reviewing on behalf of LiJunscs April 14, 2026 03:13 View session

Copilot AI reviewed Apr 14, 2026

View reviewed changes

[train]: style, fix some typo.

1f6ad03

LiJunscs force-pushed the deepseek_related branch from ff8b531 to 1f6ad03 Compare April 14, 2026 05:56

[train]: test, set golden value of engram.

3e3fe90

LiJunscs requested a review from Copilot April 14, 2026 14:12

Copilot started reviewing on behalf of LiJunscs April 14, 2026 14:13 View session

Copilot AI reviewed Apr 14, 2026

View reviewed changes

zhaoyinglia approved these changes Apr 15, 2026

View reviewed changes

lxd-cumt approved these changes Apr 15, 2026

View reviewed changes

zhaoyinglia merged commit c0b1dc5 into flagos-ai:main Apr 16, 2026
56 of 60 checks passed

		rank=get_pg_rank(self.tp_group),
		world_size=get_pg_size(self.tp_group),

		warnings.warn(f"[rank0]: We do not recomend using allreduce for engram embedding, this is deprecated and will be removed in later version.", DeprecationWarning)
		if self.args.engram_embedding_parallel_size is not None:

		assert not self.args.use_megatron_fsdp, "Megatron FSDP does not be supported yet, looking forward to later version."
		assert not self.args.init_model_with_meta_device, "Init_model_with_meta_device does not be supported yet, looking forward to later version."

	"Of cource you can set the --optimizer-offload-fraction to offload other params meanwhile enable this to offload all embedding optimizer states to CPU.")
	"Of course you can set the --optimizer-offload-fraction to offload other params meanwhile enable this to offload all embedding optimizer states to CPU.")

	metadata: Optional[dict] = None,** kwargs,
	metadata: Optional[dict] = None, **kwargs,

	# Actually, engram really causes the layers to be non-homogeneous.
	# Actually, engram really causes the layers to be non-homogeneous.
	metadata = {} if metadata is None else metadata



		class EngramMemory(nn.Module):
		"""Embedding parallelized in the vocabulay dimension.

-        # Engram let the layers be non-homogeneous, so we need to set the flag in metadata to let the sharded state dict logic know.
-        # This is usefule when all layer are same, the TransformerBlock will be homogeneous, it generate sharded_state_dict will same keys for all layer and need all layers have the same structure.
-        # The layer has engram module does not fit this assumption.
-        # If the flag is set to True, the sharded_state_dict will use layer_number to generate different keys for different layer, which is same to models has dense layer leading and moe layer following.
-        # Actually, engram really causes the layers to be non-homogeneous.
+        # Engram can make the layers non-homogeneous, so we need to set a flag in
+        # metadata to let the sharded state dict logic know.
+        # This is useful because when all layers are the same, TransformerBlock is
+        # homogeneous and generates the same sharded_state_dict keys for all layers,
+        # which requires all layers to have the same structure.
+        # Layers with an Engram module do not fit this assumption.
+        # If this flag is set to True, sharded_state_dict will use layer_number to
+        # generate different keys for different layers, similar to models with dense
+        # layers first and MoE layers following.

	engram_embedding_parallel_method: alltoall # alltoall, allreduce, offload
	engram_embedding_parallel_method: alltoall # alltoall, allreduce
	# engram_offload_embedding_optimizer_states: true # Optional; controls optimizer-state offloading.

Uh oh!

Conversation

LiJunscs commented Mar 12, 2026

PR Category

PR Types

PR Description

Uh oh!

CLAassistant commented Mar 12, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

lxd-cumt left a comment

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

LiJunscs Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

Copilot AI Apr 14, 2026

Choose a reason for hiding this comment

Uh oh!

zhaoyinglia left a comment

Choose a reason for hiding this comment

Uh oh!

lxd-cumt left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

CLAassistant commented Mar 12, 2026 •

edited

Loading