slapo.framework_dialect.deepspeed.pipeline

Classes:

WrappedTypeCode(value)	Type code for wrapped tensor for inter-GPU communication.
DeepSpeedPipeStageWrapper(mod, stage_id, ...)

Functions:

get_simple_nested_list_str(data)	A helper function that prints a nested structure without printing tensor values.
flat_and_name_tensor_list(data, name, suffix)	Flat the given list and assign the itemized name to each tensor.
encode_metadata(metadata)	Encode metadata to a string.
decode_metadata(metadata_str)	Decode metadata from a string.
flatten(outputs, device[, path, metadata, ret])	Flatten nested structure of outputs and make sure every output is a torch.Tensor.
unflatten(args, metadata)	The reverse function of 'flatten'.
analyze_tie_ranks(tie_weight_groups, topology)	Analyze the ranks of the tied weights for DeepSpeed pipeline.
deepspeed_pipe_engine(sch_metadata, ...)	DeepSpeed pipeline engine.

class slapo.framework_dialect.deepspeed.pipeline.WrappedTypeCode(value)

Type code for wrapped tensor for inter-GPU communication.

slapo.framework_dialect.deepspeed.pipeline.get_simple_nested_list_str(data)

A helper function that prints a nested structure without printing tensor values.

slapo.framework_dialect.deepspeed.pipeline.flat_and_name_tensor_list(data, name, suffix)

Flat the given list and assign the itemized name to each tensor. This is mainly used for liveness.

slapo.framework_dialect.deepspeed.pipeline.encode_metadata(metadata)

Encode metadata to a string. The metadata format is:

[<path>,<type>), (<path>,<type>), …]

After encoding: “<path>,<type>|<path>,<type>|…”

slapo.framework_dialect.deepspeed.pipeline.decode_metadata(metadata_str)

Decode metadata from a string.

slapo.framework_dialect.deepspeed.pipeline.flatten(outputs, device, path='', metadata=None, ret=None)

Flatten nested structure of outputs and make sure every output is a torch.Tensor. We maintain a metadata to restore the original structure. The metadata format is:

[<path>,<type>), (<path>,<type>), …]

• <path> is the path to the tensor in the nested structure. For example,

the paths of t1 and t2 in [[t1, t2]] are “0.0” and “0.1”. The path of t1 and t2 in [[t1], [t2]] is “0.0” and “1.0”.

• <type> is WrappedTypeCode of the tensor.

Note that len(metadata) == len(ret)

slapo.framework_dialect.deepspeed.pipeline.unflatten(args, metadata)

The reverse function of ‘flatten’.

slapo.framework_dialect.deepspeed.pipeline.analyze_tie_ranks(tie_weight_groups, topology)

Analyze the ranks of the tied weights for DeepSpeed pipeline.

class slapo.framework_dialect.deepspeed.pipeline.DeepSpeedPipeStageWrapper(mod, stage_id, name, total_stages, liveness, stage_id_2_arg_names)

Methods:

forward(*args, **kwargs)

Defines the computation performed at every call.

Parameters

• mod (fx.GraphModule) –

• stage_id (int) –

• name (str) –

• total_stages (int) –

• liveness (dict) –

• stage_id_2_arg_names (dict) –

forward(*args, **kwargs)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

slapo.framework_dialect.deepspeed.pipeline.deepspeed_pipe_engine(sch_metadata, stage_modules, **kwargs)

DeepSpeed pipeline engine.

Parameters

• sch_metadata (ScheduleMetadata) – The schedule metadata.

• stage_modules (List[nn.Module]) – The list of pipeline stage modules.

• **kwargs – The keyword arguments. Should include DeepSpeed related information, such as “config”, “loss_fn”, “topology”, “fp16”.

Returns

model – The DeepSpeed pipeline module.

Return type

PipelineModule