TRN2 Meshes and Configurations #916

apoorvtintin · 2025-01-10T00:48:35Z

This PR adds meshes for TRN2/1 for Fuji models and transformer layer configuration favorable to Neuron.

Neuron supports stacked transformer and GroupedQKVLinear instead of FusedGroupedQKVLinear for Grouped Query Attention (GQA)

This is a newer version of the PR #885. This PR resolved all comments and requested changes mentioned in the linked PR.

apoorvtintin · 2025-01-10T00:55:07Z

Added a ModelConfigModifier that overrides the class for a module. Allowing different model configurations based on Model size and platform.

kelvin-zou

Thank you for making such change, overall looks good. A few nit comments.

kelvin-zou · 2025-01-10T00:57:05Z

axlearn/common/trainer_config_modifier.py

+                continue
+            # Here we assume x.y.z format.
+            # One example would be model.decoder.transformer.layer.
+            target_modules = module_name.split(".")


Can you try to extract a common util function named something like
def replace_module_recursive(target_modules:str, config_key: str, target_config) and make it applied to both here and RematSpecModifier

I extracted a helper function, let me know if this looks good

axlearn/common/trainer_config_modifier_test.py

apoorvtintin · 2025-01-10T07:44:33Z

Added ParameterPartitionSpecModifier for parameters to shard Embeddings in a vocab parallel manner as described in Megatron LM.

axlearn/common/trainer_config_modifier.py

ruomingp · 2025-01-11T23:50:52Z

axlearn/common/trainer_config_modifier.py

+
+            found_module, parent_module, key_in_parent = find_target_module(module_name, cfg)
+
+            # Copy configurations from the config being replaced on a best effort basis


Wait, this behavior is not explained in the class comments. So we are not replacing but merging the configs? Maybe we should support a merge function instead?

Yeah the goal is to change the config to a similar module. This means most of the configuration can be reused from before. Essentially replacing the module but merging the config. Let me extract out a merge function.

Abstracted out a merge function let me know if more changes are needed for this.

apoorvtintin · 2025-01-12T07:08:47Z

@ruomingp Thank you for the review, I have addressed all your comments, please let me know if more changes are needed.

axlearn/common/trainer_config_modifier.py

ruomingp · 2025-01-13T06:35:53Z

axlearn/common/trainer_config_modifier.py

+        for module_name, model_cfg in self._model_cfg_modifications.items():
+            found_module = _find_target_module(module_name, cfg)


In utils.py we have get_recursively and set_recursively for Nested[...]. I wonder if it will be useful to add corresponding methods to ConfigBase. Then we can do something like:

Suggested change

for module_name, model_cfg in self._model_cfg_modifications.items():

found_module = _find_target_module(module_name, cfg)

for cfg_path, cfg_modification in self._model_cfg_modifications.items():

child_cfg = cfg.get_recursively(cfg_path)

child_cfg = cfg_modification(child_cfg, path=cfg_path)

cfg.set_recursively(cfg_path, value=child_cfg)

Added get_recursively and set_recursively functions to ConfigBase. Let me know if it looks good

I wonder if an alternative (which aims to simplify the ConfigBase api) is to do something similar to Python's sorted; we allow utils.get_resursively to take a value fn:

# Default behavior is to use key lookup: utils.get_recursively(..., value_fn=lambda k,v: v[k]) # Custom behavior can be attribute lookup: utils.get_recursively(..., value_fn=lambda k,v: getattr(v,k))

A benefit is that other non-config instances can also leverage get_recursively.

Hi @markblee , maybe we can do this in a follow-up PR?

apoorvtintin · 2025-01-15T00:26:21Z

Added a more flexible PartitionSpecModifier that can modify multiple partition_spec attributes in a single module config.

kelvin-zou

Mostly lgtm, some minor comments.

axlearn/common/trainer_config_modifier.py

apoorvtintin · 2025-01-27T10:09:17Z

@ruomingp and @kelvin-zou thank you both for the review. I addressed all comments, please let me know if anymore changes are needed. PR looks clean now.

markblee · 2025-01-27T17:46:24Z

axlearn/common/config.py

+        key: str
+
+    def recursive_traverse(self, key_path: Sequence[str]) -> tuple[Any, str]:
+        """Recursively traverse the config to find the target key.


Please see other comment re get_recursively; also, I wonder whether we actually need recursion here (seems like a loop would be simpler).

Both loops and recursion are fair choices I was following the pattern of utils.get_recusively() based on previous guidance.

markblee · 2025-01-27T17:47:21Z

axlearn/common/trainer_config_modifier.py

@@ -146,6 +137,110 @@ def __call__(self, cfg: SpmdTrainer.Config) -> SpmdTrainer.Config:
        return cfg


+class ModelConfigModifier(ConfigModifier):


Which part of this class is specific to model? It seems to take generic modifications?

The class makes some assumptions around the klass attribute of configs when merging/reusing the previous configs which makes it non generic.

@markblee I would love to get some more guidance on a name that fits this class better, do you have any suggestions?

Hm, but all configs deriving from Configurable have klass. Did you mean ModuleConfigModifier perhaps, as it replaces arbitrary modules by path? (model -> module?)

Yeah I agree ModuleConfigModifier is better, thanks!

axlearn/common/trainer_config_modifier.py

markblee · 2025-01-27T17:54:03Z

axlearn/experiments/text/gpt/fuji.py

+    if version != Version.V1:
+        trn2_model_modifications.append(
+            ModelConfigModifier.default_config().set(
+                target_config="model.decoder.transformer.layer.self_attention.attention."


A downside of representing these deeply nested configs as string paths is that they are brittle, and can quickly become outdated.

Have we considered using cfg.visit to achieve some of these modifications (e.g.,

axlearn/axlearn/common/layers.py

Line 266 in 1c22688

def set_layer_norm_eps_recursively(cfg: ConfigBase, eps: float, set_only_if_none: bool = False):

)?

(A bit late to review, so apologies if this discussion has already taken place.)

axlearn/common/config.py

ruomingp · 2025-01-27T19:10:40Z

axlearn/common/config.py

+        """Recursively find the target key in the config and return its value.
+
+        Args:
+            key_path: A sequence of keys for indexing to get the target value.


Can path be empty? Maybe it can return self if path is empty?

Added a check to see if path is empty, throws value error now if path is empty

axlearn/common/config.py

ruomingp · 2025-01-27T19:12:21Z

axlearn/common/config.py

+        Raises:
+            ValueError: A key in key_path is not found.
+        """
+        traverse_result = self.recursive_traverse(key_path)


Can we do something like:

Suggested change

traverse_result = self.recursive_traverse(key_path)

if not path:

raise ValueError(...)

parent = self.get_recursively(path[:-1])

I removed the recursive_traverse method and added a value error if path is empty. I am still using path[1:] instead of path[:-1] since we are traversing parents first.

apoorvtintin · 2025-01-31T19:21:46Z

Thank you for the reviews @markblee @ruomingp @kelvin-zou!
@ruomingp I addressed most of your comments, please let me know if more changes are needed. @markblee We can discuss the usage of cfg.visit instead of using full module paths further, but I would like to address this in a follow up PR if possible.

axlearn/common/config.py

markblee

Hi @apoorvtintin, please feel free leave TODO(markblee): for my comments and go ahead once @ruomingp approves. I'll be happy to send a follow-up PR to avoid blocking this one.

markblee · 2025-02-01T02:08:43Z

axlearn/common/trainer_config_modifier.py

@@ -146,6 +137,110 @@ def __call__(self, cfg: SpmdTrainer.Config) -> SpmdTrainer.Config:
        return cfg


+class ModelConfigModifier(ConfigModifier):


Hm, but all configs deriving from Configurable have klass. Did you mean ModuleConfigModifier perhaps, as it replaces arbitrary modules by path? (model -> module?)

axlearn/experiments/text/gpt/fuji.py

apoorvtintin · 2025-02-03T09:29:56Z

Addressed all comments. @markblee I added TODOs for us to revisit cfg.visit later, please let me know if I should add more.

ruomingp · 2025-02-03T17:17:39Z

axlearn/common/config.py

+
+            if index == len(path):
+                return current


Suggested change

if index == len(path):

return current

ruomingp · 2025-02-03T17:19:11Z

axlearn/common/config.py

+        current = self
+        for i, key in enumerate(path):
+            if i == len(path) - 1:
+                setattr(current, key, value)
+                return
+            else:
+                # TODO(markblee): maybe use cfg.visit instead of getattr
+                current = getattr(current, key)


Suggested change

current = self

for i, key in enumerate(path):

if i == len(path) - 1:

setattr(current, key, value)

return

else:

# TODO(markblee): maybe use cfg.visit instead of getattr

current = getattr(current, key)

parent = self.get_recursively(path[:-1])

setattr(current, path[-1], value)

+ address comments

apoorvtintin requested review from ruomingp, markblee and a team as code owners January 10, 2025 00:48

apoorvtintin mentioned this pull request Jan 10, 2025

Add meshes and config for TRN2/1 for Fuji models #885

Closed

apoorvtintin force-pushed the mainline-upstream-boilerplate branch 2 times, most recently from 6b404f6 to 3f7c840 Compare January 10, 2025 00:53

kelvin-zou reviewed Jan 10, 2025

View reviewed changes

apoorvtintin force-pushed the mainline-upstream-boilerplate branch 2 times, most recently from 708fc5e to d481132 Compare January 10, 2025 07:38

apoorvtintin force-pushed the mainline-upstream-boilerplate branch 2 times, most recently from 5be50d7 to 9b10041 Compare January 10, 2025 08:10

apoorvtintin requested a review from kelvin-zou January 10, 2025 19:44

ruomingp requested changes Jan 11, 2025

View reviewed changes

apoorvtintin force-pushed the mainline-upstream-boilerplate branch from 9b10041 to 0f0a530 Compare January 12, 2025 07:06

apoorvtintin requested a review from ruomingp January 12, 2025 07:08

ruomingp requested changes Jan 13, 2025

View reviewed changes

apoorvtintin mentioned this pull request Jan 13, 2025

[DO-NOT-MERGE] PR encompassing all changes needed to support neuron on Axlearn #919

Open

apoorvtintin force-pushed the mainline-upstream-boilerplate branch 2 times, most recently from c23e3b2 to 94bfff6 Compare January 15, 2025 00:10

apoorvtintin force-pushed the mainline-upstream-boilerplate branch 2 times, most recently from 45c7df1 to 8807856 Compare January 17, 2025 01:17

kelvin-zou reviewed Jan 17, 2025

View reviewed changes

axlearn/common/trainer_config_modifier.py Outdated Show resolved Hide resolved

axlearn/common/trainer_config_modifier.py Outdated Show resolved Hide resolved

apoorvtintin force-pushed the mainline-upstream-boilerplate branch from 8807856 to 25510d6 Compare January 22, 2025 01:39

apoorvtintin requested a review from kelvin-zou January 22, 2025 01:40

apoorvtintin force-pushed the mainline-upstream-boilerplate branch 2 times, most recently from eec33eb to 86bafa8 Compare January 23, 2025 05:40

apoorvtintin force-pushed the mainline-upstream-boilerplate branch 3 times, most recently from fe96240 to da90757 Compare January 24, 2025 23:32

apoorvtintin requested a review from ruomingp January 24, 2025 23:40

apoorvtintin force-pushed the mainline-upstream-boilerplate branch from da90757 to 7e2e5f2 Compare January 27, 2025 10:07

markblee reviewed Jan 27, 2025

View reviewed changes

ruomingp reviewed Jan 27, 2025

View reviewed changes

axlearn/common/config.py Outdated Show resolved Hide resolved

axlearn/common/config.py Outdated Show resolved Hide resolved

ruomingp reviewed Jan 27, 2025

View reviewed changes

apoorvtintin added 2 commits January 31, 2025 11:03

TRN2 Meshes and Configurations

17c7d66

Add get_recursive and set_recursive to ConfigBase.

71e006a

apoorvtintin force-pushed the mainline-upstream-boilerplate branch from 7e2e5f2 to 73e14a6 Compare January 31, 2025 19:05

apoorvtintin requested review from ruomingp and markblee January 31, 2025 19:18

ruomingp reviewed Jan 31, 2025

View reviewed changes

markblee reviewed Feb 1, 2025

View reviewed changes

apoorvtintin force-pushed the mainline-upstream-boilerplate branch from 73e14a6 to 61efbb5 Compare February 3, 2025 09:28

apoorvtintin requested review from ruomingp and markblee February 3, 2025 09:28

ruomingp reviewed Feb 3, 2025

View reviewed changes

apoorvtintin added 2 commits February 3, 2025 09:45

Use loops inside get/set_recursively

8d976db

+ address comments

Update golden configs

dfbb412

apoorvtintin force-pushed the mainline-upstream-boilerplate branch from 61efbb5 to dfbb412 Compare February 3, 2025 18:45

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

TRN2 Meshes and Configurations #916

TRN2 Meshes and Configurations #916

apoorvtintin commented Jan 10, 2025

apoorvtintin commented Jan 10, 2025

kelvin-zou left a comment

kelvin-zou Jan 10, 2025

apoorvtintin Jan 10, 2025 •

edited

Loading

apoorvtintin commented Jan 10, 2025 •

edited

Loading

ruomingp Jan 11, 2025

apoorvtintin Jan 12, 2025

apoorvtintin Jan 12, 2025

apoorvtintin commented Jan 12, 2025

ruomingp Jan 13, 2025

apoorvtintin Jan 27, 2025

markblee Jan 27, 2025

ruomingp Jan 31, 2025

apoorvtintin commented Jan 15, 2025 •

edited

Loading

kelvin-zou left a comment

apoorvtintin commented Jan 27, 2025 •

edited

Loading

markblee Jan 27, 2025

apoorvtintin Jan 31, 2025

markblee Jan 27, 2025

apoorvtintin Jan 31, 2025 •

edited

Loading

apoorvtintin Jan 31, 2025

markblee Feb 1, 2025

apoorvtintin Feb 3, 2025

markblee Jan 27, 2025

ruomingp Jan 27, 2025

apoorvtintin Jan 31, 2025

ruomingp Jan 27, 2025

apoorvtintin Jan 31, 2025 •

edited

Loading

apoorvtintin commented Jan 31, 2025

markblee left a comment

markblee Feb 1, 2025

apoorvtintin commented Feb 3, 2025 •

edited

Loading

ruomingp Feb 3, 2025

ruomingp Feb 3, 2025


		found_module, parent_module, key_in_parent = find_target_module(module_name, cfg)

		# Copy configurations from the config being replaced on a best effort basis

		for module_name, model_cfg in self._model_cfg_modifications.items():
		found_module = _find_target_module(module_name, cfg)

-        for module_name, model_cfg in self._model_cfg_modifications.items():
-            found_module = _find_target_module(module_name, cfg)
+        for cfg_path, cfg_modification in self._model_cfg_modifications.items():
+            child_cfg = cfg.get_recursively(cfg_path)
+            child_cfg = cfg_modification(child_cfg, path=cfg_path)
+            cfg.set_recursively(cfg_path, value=child_cfg)

		@@ -146,6 +137,110 @@ def __call__(self, cfg: SpmdTrainer.Config) -> SpmdTrainer.Config:
		return cfg


		class ModelConfigModifier(ConfigModifier):

-        traverse_result = self.recursive_traverse(key_path)
+        if not path:
+            raise ValueError(...)
+        parent = self.get_recursively(path[:-1])

TRN2 Meshes and Configurations #916

Are you sure you want to change the base?

TRN2 Meshes and Configurations #916

Conversation

apoorvtintin commented Jan 10, 2025

apoorvtintin commented Jan 10, 2025

kelvin-zou left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

apoorvtintin Jan 10, 2025 • edited Loading

Choose a reason for hiding this comment

apoorvtintin commented Jan 10, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

apoorvtintin commented Jan 12, 2025

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

apoorvtintin commented Jan 15, 2025 • edited Loading

kelvin-zou left a comment

Choose a reason for hiding this comment

apoorvtintin commented Jan 27, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

apoorvtintin Jan 31, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

apoorvtintin Jan 31, 2025 • edited Loading

Choose a reason for hiding this comment

apoorvtintin commented Jan 31, 2025

markblee left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

apoorvtintin commented Feb 3, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

apoorvtintin Jan 10, 2025 •

edited

Loading

apoorvtintin commented Jan 10, 2025 •

edited

Loading

apoorvtintin commented Jan 15, 2025 •

edited

Loading

apoorvtintin commented Jan 27, 2025 •

edited

Loading

apoorvtintin Jan 31, 2025 •

edited

Loading

apoorvtintin Jan 31, 2025 •

edited

Loading

apoorvtintin commented Feb 3, 2025 •

edited

Loading