module
module
¶
A theseus nn module, which is nn module, execpt it asks you to dump all the types of constitutent parts so we can configure it.
Module
¶
Bases: Module
sharding: ShardingPlan
property
¶
Logical mappings; models without a plan use replicated parameters.
flops(seq: int) -> float
¶
Estimated forward + backward FLOPs for one sequence of length seq.
Composite modules require binding to existing variables so their actual setup children are available. Counts use dense matmul arithmetic (one multiply-add is two FLOPs, backward is twice forward), excluding optimizer updates, rematerialization and minor elementwise operations. Unsupported operations contribute zero; this is not an exact counter.
components() -> List[Type[Any]]
classmethod
¶
Return the types of constituent parts of this module.
Returns:
| Name | Type | Description |
|---|---|---|
Type |
List[Type[Any]]
|
A type or tuple of types representing the constituent parts. |
gather() -> List[Type[Any]]
classmethod
¶
Depth-first search of all constituent parts of this module.
Returns:
| Type | Description |
|---|---|
List[Type[Any]]
|
List[Type]: A list of all constituent part types. |